feat(link): unified AI login-manager registry (/v1/links)

The org+user-scoped registry of which provider accounts (Claude Max, ChatGPT
Plus, a Hanzo/api key) a developer has signed into, on which machines, with each
account's latest usage snapshot — the cross-machine view console renders and the
source the redundancy route policy reads.

- clients/link: the Link atom (no secret — metadata + usage snapshot only),
  per-org SQLite store (org+subject leading-bound, upsert-on-identity, revoke),
  the /v1/links surface, and a pure RoutePolicy (Plan) that orders a user's linked
  accounts for redundancy (two Claude Max, then the metered API backstop) carrying
  the billing mode per candidate. The store holds no metering client — a
  subscription's usage is metered for visibility only and never charges commerce.
- clients/agents: a session now carries the linked account it ran under
  (Provider/Account tag), and StopSessions/CountActiveSessions expose the
  in-process action a link revoke takes to stop the sessions under a revoked
  account/device. Backward-compatible session migration (addColumns).
- subsystems: mount link after agents so a revoke can stop its sessions.

Org+user fail-closed isolation, subscription-vs-api-key billing distinction,
and revoke-stops-sessions are all tested (build/vet/gofmt clean; race+CGO green).
This commit is contained in:
2026-07-15 03:31:37 -07:00
parent bc328bdb5d
commit a89c7b08c0
13 changed files with 1860 additions and 14 deletions
+21 -4
View File
@@ -61,6 +61,8 @@ const (
maxHost = 256
maxCwd = 1024
maxRepo = 512
maxProvider = 64
maxAccount = 256
)
func validKind(k string) bool {
@@ -85,10 +87,12 @@ type sessionView struct {
TaskRunID string `json:"taskRunId,omitempty"`
// Execution context (mission-control): the machine/repo/cwd a card shows and
// the run-target a session is dispatched to. Omitted when a surface didn't report it.
Host string `json:"host,omitempty"`
Cwd string `json:"cwd,omitempty"`
Repo string `json:"repo,omitempty"`
Target string `json:"target,omitempty"`
Host string `json:"host,omitempty"`
Cwd string `json:"cwd,omitempty"`
Repo string `json:"repo,omitempty"`
Target string `json:"target,omitempty"`
Provider string `json:"provider,omitempty"`
Account string `json:"account,omitempty"`
Events int `json:"events"`
Children int `json:"children"`
@@ -155,6 +159,7 @@ func toSessionView(x Session, events, children int) sessionView {
ParentSessionID: x.ParentID, RootSessionID: x.RootID, Title: x.Title,
TaskWorkflowID: x.TaskWorkflowID, TaskRunID: x.TaskRunID,
Host: x.Host, Cwd: x.Cwd, Repo: x.Repo, Target: x.Target,
Provider: x.Provider, Account: x.Account,
Events: events, Children: children,
StartedAt: rfc3339(x.StartedAt), EndedAt: rfc3339(x.EndedAt),
CreatedAt: rfc3339(x.CreatedAt), UpdatedAt: rfc3339(x.UpdatedAt),
@@ -207,6 +212,9 @@ type registerReq struct {
Cwd string `json:"cwd"`
Repo string `json:"repo"`
Target string `json:"target"`
// Account tag — the linked AI account this session ran under (login manager).
Provider string `json:"provider"`
Account string `json:"account"`
}
func registerSession(s *cloud.Service[state], c *zip.Ctx) error {
@@ -249,6 +257,14 @@ func registerSession(s *cloud.Service[state], c *zip.Ctx) error {
if cerr != nil {
return cerr
}
provider := strings.TrimSpace(body.Provider)
account := strings.TrimSpace(body.Account)
if len(provider) > maxProvider {
return zip.ErrBadRequest("provider too long")
}
if len(account) > maxAccount {
return zip.ErrBadRequest("account too long")
}
id, err := genID("sess")
if err != nil {
@@ -261,6 +277,7 @@ func registerSession(s *cloud.Service[state], c *zip.Ctx) error {
TaskWorkflowID: strings.TrimSpace(body.TaskWorkflowID),
TaskRunID: strings.TrimSpace(body.TaskRunID),
Host: host, Cwd: cwd, Repo: repo, Target: target,
Provider: provider, Account: account,
StartedAt: now, CreatedAt: now, UpdatedAt: now,
}
if isTerminalStatus(status) {
+157
View File
@@ -0,0 +1,157 @@
package agents
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
)
// sessions_stop.go is the login-manager tie-in: the in-process action a link
// revoke takes to tear down the live sessions that ran under a revoked account or
// device, plus the active-session count the device view shows. Both are org-scoped
// (org is the ONLY tenant key) and nil-safe (no agents mounted → 0), and neither
// can fan out to another tenant or to an org's every session by accident.
// SessionMatch selects live (running|paused) sessions to stop or count. Fields are
// ANDed with the org; an empty field is "any". Host targets a device; Provider +
// Account target a linked AI account. An all-empty match selects NOTHING, so a
// stop can never sweep an org's every session by accident.
type SessionMatch struct {
Host string
Provider string
Account string
}
func (m SessionMatch) empty() bool {
return strings.TrimSpace(m.Host) == "" &&
strings.TrimSpace(m.Provider) == "" &&
strings.TrimSpace(m.Account) == ""
}
// where builds the ANDed predicate + args for a live-session match under org.
func (m SessionMatch) where(org string) (string, []any) {
where := "org=? AND status IN (?,?)"
args := []any{org, StatusRunning, StatusPaused}
if h := strings.TrimSpace(m.Host); h != "" {
where += " AND host=?"
args = append(args, h)
}
if p := strings.TrimSpace(m.Provider); p != "" {
where += " AND provider=?"
args = append(args, p)
}
if a := strings.TrimSpace(m.Account); a != "" {
where += " AND account=?"
args = append(args, a)
}
return where, args
}
// listActiveMatch returns org's live sessions matching m, oldest first.
func (s *Store) listActiveMatch(ctx context.Context, org string, m SessionMatch) ([]Session, error) {
where, args := m.where(org)
rows, err := s.db.QueryContext(ctx,
`SELECT `+sessionCols+` FROM agent_sessions WHERE `+where+` ORDER BY created_at ASC`, args...)
if err != nil {
return nil, fmt.Errorf("list active match: %w", err)
}
defer func() { _ = rows.Close() }()
var out []Session
for rows.Next() {
x, err := scanSession(rows)
if err != nil {
return nil, fmt.Errorf("scan session: %w", err)
}
out = append(out, x)
}
return out, rows.Err()
}
// countActiveMatch counts org's live sessions matching m.
func (s *Store) countActiveMatch(ctx context.Context, org string, m SessionMatch) (int, error) {
where, args := m.where(org)
var n int
if err := s.db.QueryRowContext(ctx,
`SELECT COUNT(*) FROM agent_sessions WHERE `+where, args...).Scan(&n); err != nil {
return 0, fmt.Errorf("count active match: %w", err)
}
return n, nil
}
// StopSessions closes every RUNNING|PAUSED session of org matching m — recording a
// control "stop" event on each and transitioning it to a terminal state — and
// returns how many it stopped. It is the action a login-out (link revoke) takes so
// the sessions that ran under a revoked account/device are torn down. Org is the
// ONLY tenant key; a caller for org A can never stop org B's sessions. An all-empty
// match stops nothing (never an accidental org-wide stop). Not-mounted → (0, nil),
// so a revoke tolerates a deployment with no session plane.
func StopSessions(ctx context.Context, org string, m SessionMatch) (int, error) {
if mounted == nil {
return 0, nil
}
org = strings.TrimSpace(org)
if org == "" || m.empty() {
return 0, nil
}
live, err := mounted.State.store.listActiveMatch(ctx, org, m)
if err != nil {
return 0, err
}
stopped := 0
for _, x := range live {
if err := stopOne(ctx, x); err != nil {
// Best-effort per session: a failure on one does not abort the rest, so a
// revoke tears down as many as it can and reports the true count.
mounted.Log.Warn("agents: stop session", "org", org, "session", x.ID, "err", err)
continue
}
stopped++
}
return stopped, nil
}
// stopOne records a stop control event on a live session and moves it to a
// terminal (error) state — the forced-teardown transition. A session already
// terminal is skipped (monotonic terminal rule).
func stopOne(ctx context.Context, x Session) error {
if isTerminalStatus(x.Status) {
return nil
}
now := time.Now().Unix()
if evID, err := genID("evt"); err == nil {
payload, _ := json.Marshal(controlPayload{Command: CmdStop, Message: "account logged out via login manager"})
e, aerr := mounted.State.store.AppendEvent(ctx, Event{
ID: evID, SessionID: x.ID, Org: x.Org, Kind: KindControl,
Actor: billingActor(x.Org, ""), Payload: string(payload), CreatedAt: now,
})
if aerr == nil {
publishEvent(mounted, x.Org, x.RootID, e)
}
}
x.Status = StatusError
x.EndedAt = now
x.UpdatedAt = now
if err := mounted.State.store.UpdateSession(ctx, x); err != nil {
return err
}
ev, _ := mounted.State.store.CountEvents(ctx, x.Org, x.ID)
ch, _ := mounted.State.store.CountChildren(ctx, x.Org, x.ID)
publishSession(mounted, x, ev, ch)
return nil
}
// CountActiveSessions returns how many of org's sessions matching m are live
// (running|paused) — the device view's "active sessions". Org-scoped; 0 when not
// mounted or the match is empty.
func CountActiveSessions(ctx context.Context, org string, m SessionMatch) (int, error) {
if mounted == nil {
return 0, nil
}
org = strings.TrimSpace(org)
if org == "" || m.empty() {
return 0, nil
}
return mounted.State.store.countActiveMatch(ctx, org, m)
}
+24 -10
View File
@@ -52,6 +52,14 @@ type Session struct {
Cwd string
Repo string
Target string
// Provider/Account tag a session with the linked AI account it ran under (the
// login-manager tie-in): which provider (claude|codex|hanzo|…) and which
// subscription/api account served this run. Optional (a surface that doesn't
// know sets ""), surfaced so the cockpit shows "this ran on your Claude Max
// acct" and so a login-out (link revoke) can stop the sessions that used it.
Provider string
Account string
}
// Event is one entry in a session's ordered log: a model message, a tool call, a
@@ -116,7 +124,9 @@ CREATE TABLE IF NOT EXISTS agent_sessions (
host TEXT NOT NULL DEFAULT '',
cwd TEXT NOT NULL DEFAULT '',
repo TEXT NOT NULL DEFAULT '',
target TEXT NOT NULL DEFAULT ''
target TEXT NOT NULL DEFAULT '',
provider TEXT NOT NULL DEFAULT '',
account TEXT NOT NULL DEFAULT ''
);
CREATE INDEX IF NOT EXISTS ix_sessions_org_root ON agent_sessions(org, root_id, created_at);
CREATE INDEX IF NOT EXISTS ix_sessions_org_parent ON agent_sessions(org, parent_id, created_at);
@@ -141,10 +151,12 @@ CREATE INDEX IF NOT EXISTS ix_events_org_session_seq ON agent_session_events(org
// Forward, idempotent: a sessions table created before the execution-context
// columns existed gains them here (the CREATE above only runs on a fresh DB).
if err := s.addColumns("agent_sessions", map[string]string{
"host": "TEXT NOT NULL DEFAULT ''",
"cwd": "TEXT NOT NULL DEFAULT ''",
"repo": "TEXT NOT NULL DEFAULT ''",
"target": "TEXT NOT NULL DEFAULT ''",
"host": "TEXT NOT NULL DEFAULT ''",
"cwd": "TEXT NOT NULL DEFAULT ''",
"repo": "TEXT NOT NULL DEFAULT ''",
"target": "TEXT NOT NULL DEFAULT ''",
"provider": "TEXT NOT NULL DEFAULT ''",
"account": "TEXT NOT NULL DEFAULT ''",
}); err != nil {
return err
}
@@ -154,19 +166,21 @@ CREATE INDEX IF NOT EXISTS ix_events_org_session_seq ON agent_session_events(org
// column and fail on the old schema ("no such column: target").
if _, err := s.db.Exec(`
CREATE INDEX IF NOT EXISTS ix_sessions_org_target ON agent_sessions(org, target);
CREATE INDEX IF NOT EXISTS ix_sessions_org_host ON agent_sessions(org, host);`); err != nil {
CREATE INDEX IF NOT EXISTS ix_sessions_org_host ON agent_sessions(org, host);
CREATE INDEX IF NOT EXISTS ix_sessions_org_account ON agent_sessions(org, provider, account);`); err != nil {
return fmt.Errorf("migrate sessions indexes: %w", err)
}
return nil
}
const sessionCols = `id,org,agent,actor,status,parent_id,root_id,title,started_at,ended_at,created_at,updated_at,task_workflow_id,task_run_id,host,cwd,repo,target`
const sessionCols = `id,org,agent,actor,status,parent_id,root_id,title,started_at,ended_at,created_at,updated_at,task_workflow_id,task_run_id,host,cwd,repo,target,provider,account`
func scanSession(sc interface{ Scan(...any) error }) (Session, error) {
var x Session
err := sc.Scan(&x.ID, &x.Org, &x.Agent, &x.Actor, &x.Status, &x.ParentID, &x.RootID,
&x.Title, &x.StartedAt, &x.EndedAt, &x.CreatedAt, &x.UpdatedAt,
&x.TaskWorkflowID, &x.TaskRunID, &x.Host, &x.Cwd, &x.Repo, &x.Target)
&x.TaskWorkflowID, &x.TaskRunID, &x.Host, &x.Cwd, &x.Repo, &x.Target,
&x.Provider, &x.Account)
return x, err
}
@@ -191,10 +205,10 @@ func (s *Store) CreateSession(ctx context.Context, x Session) error {
}
}
_, err := s.db.ExecContext(ctx,
`INSERT INTO agent_sessions (`+sessionCols+`) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
`INSERT INTO agent_sessions (`+sessionCols+`) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
x.ID, x.Org, x.Agent, x.Actor, x.Status, x.ParentID, x.RootID, x.Title,
x.StartedAt, x.EndedAt, x.CreatedAt, x.UpdatedAt, x.TaskWorkflowID, x.TaskRunID,
x.Host, x.Cwd, x.Repo, x.Target)
x.Host, x.Cwd, x.Repo, x.Target, x.Provider, x.Account)
if err != nil {
return fmt.Errorf("insert session: %w", err)
}
+27
View File
@@ -0,0 +1,27 @@
package link
import (
"context"
"github.com/hanzoai/cloud/clients/agents"
)
// adapters.go binds the Sessions seam to the agents in-process control plane
// (clients/agents). It is the ONLY file in clients/link that imports agents;
// http.go/store.go/route.go stay free of it so the orchestration is unit-tested
// against a fake seam. agents does NOT import link, so this direction is
// cycle-free.
type sessionAdapter struct{}
func (sessionAdapter) Stop(ctx context.Context, org string, m SessionMatch) (int, error) {
return agents.StopSessions(ctx, org, agents.SessionMatch{
Host: m.Host, Provider: m.Provider, Account: m.Account,
})
}
func (sessionAdapter) CountActive(ctx context.Context, org string, m SessionMatch) (int, error) {
return agents.CountActiveSessions(ctx, org, agents.SessionMatch{
Host: m.Host, Provider: m.Provider, Account: m.Account,
})
}
+425
View File
@@ -0,0 +1,425 @@
package link
import (
"context"
"crypto/rand"
"encoding/hex"
"encoding/json"
"fmt"
"net/http"
"os"
"path/filepath"
"time"
"github.com/hanzoai/cloud"
"github.com/hanzoai/cloud/clients/principal"
"github.com/zap-proto/zip"
)
// This file mounts the login-manager registry under /v1/links — the cross-machine
// account view console renders and the collector reports into.
//
// POST /v1/links register/upsert a link (device+account+usage) -> Link
// GET /v1/links the caller's links + device projection -> {links,devices}
// GET /v1/links/route the redundancy route plan across the caller's accounts -> RoutePlan
// GET /v1/links/devices/:machine one device: its accounts + usage + active sessions -> Device
// POST /v1/links/devices/:machine/revoke revoke every account on a device + stop its sessions
// GET /v1/links/:id one link -> Link
// DELETE /v1/links/:id revoke a link (log out) + stop its sessions
//
// Every route is org+user scoped through principal.Org (a validated principal AND
// a non-empty org) plus c.User() (the owning subject), so a caller sees and
// mutates only their OWN accounts — cross-tenant and cross-user access is refused
// fail-closed.
// SessionMatch selects the live sessions a revoke stops. Fields are ANDed with the
// org; an empty field is "any". A link revoke matches {Host,Provider,Account} (the
// device+account the sessions ran under); a device revoke matches {Host}.
type SessionMatch struct {
Host string
Provider string
Account string
}
// Sessions is the seam to the agent-session control plane (clients/agents,
// in-process). Revoke stops the sessions that ran under a revoked account/device;
// the device view counts a machine's active sessions. A nil seam (unit test / no
// agents mounted) makes revoke skip the stop and the count report 0 — the registry
// truth (the revoked row) is unaffected.
type Sessions interface {
Stop(ctx context.Context, org string, m SessionMatch) (int, error)
CountActive(ctx context.Context, org string, m SessionMatch) (int, error)
}
type state struct {
store *Store
sessions Sessions
}
var mounted *cloud.Service[state]
// Mount wires the /v1/links surface. The sessions seam is set from the agents
// in-process adapter (adapters.go) so a revoke can stop the affected sessions.
func Mount(app *zip.App, deps cloud.Deps) error {
if app == nil {
return fmt.Errorf("link.Mount: nil zip.App")
}
if deps.Logger == nil {
return fmt.Errorf("link.Mount: nil deps.Logger")
}
log := deps.Logger.New("subsystem", "link")
if deps.DataDir == "" {
return fmt.Errorf("link.Mount: empty DataDir")
}
if err := os.MkdirAll(deps.DataDir, 0o755); err != nil {
return fmt.Errorf("link.Mount: data dir: %w", err)
}
store, err := openStore(filepath.Join(deps.DataDir, "link.db"))
if err != nil {
return fmt.Errorf("link.Mount: open store: %w", err)
}
s := &cloud.Service[state]{
Base: cloud.NewBase(deps, "link"),
State: state{store: store, sessions: sessionAdapter{}},
}
mounted = s
app.Post("/v1/links", cloud.Handle(s, upsertLink))
app.Get("/v1/links", cloud.Handle(s, listLinks))
// Static literals before the :id param — Fiber matches in registration order,
// so "route"/"devices" must win over :id.
app.Get("/v1/links/route", cloud.Handle(s, routePlan))
app.Get("/v1/links/devices/:machine", cloud.Handle(s, deviceDetail))
app.Post("/v1/links/devices/:machine/revoke", cloud.Handle(s, revokeDevice))
app.Get("/v1/links/:id", cloud.Handle(s, getLink))
app.Delete("/v1/links/:id", cloud.Handle(s, revokeLink))
log.Info("link mounted", "brand", deps.Brand)
return nil
}
// Shutdown closes the store. Idempotent.
func Shutdown(context.Context) error {
if mounted == nil {
return nil
}
err := mounted.State.store.Close()
mounted = nil
return err
}
// caller resolves the (org, subject) scope for a request: a validated principal
// with a non-empty org. Every handler gates on it — an off-gateway forge with no
// validated user is refused. c.User() is guaranteed non-empty once principal.Org
// returns ok (Org composes Validated, which is c.User() != "").
func caller(c *zip.Ctx) (org, user string, ok bool) {
org, ok = principal.Org(c)
if !ok {
return "", "", false
}
return org, trim(c.User()), true
}
// ---- views ----
type linkView struct {
ID string `json:"id"`
User string `json:"user"`
Machine string `json:"machine"`
Host string `json:"host,omitempty"`
OS string `json:"os,omitempty"`
Provider string `json:"provider"`
Account string `json:"account,omitempty"`
Plan string `json:"plan,omitempty"`
Kind string `json:"kind"`
Billing string `json:"billing"` // BillingMode(Kind) — how this account's usage bills
Status string `json:"status"`
LastSeen string `json:"lastSeen,omitempty"`
Usage json.RawMessage `json:"usage,omitempty"`
CreatedAt string `json:"createdAt"`
UpdatedAt string `json:"updatedAt"`
}
type deviceView struct {
Machine string `json:"machine"`
Host string `json:"host,omitempty"`
OS string `json:"os,omitempty"`
LastSeen string `json:"lastSeen,omitempty"`
Accounts []linkView `json:"accounts"`
ActiveSessions int `json:"activeSessions"`
}
func rfc3339(unix int64) string {
if unix == 0 {
return ""
}
return time.Unix(unix, 0).UTC().Format(time.RFC3339)
}
func toLinkView(l Link) linkView {
var u json.RawMessage
if l.Usage != "" {
u = json.RawMessage(l.Usage)
}
return linkView{
ID: l.ID, User: l.User, Machine: l.Machine, Host: l.Host, OS: l.OS,
Provider: l.Provider, Account: l.Account, Plan: l.Plan, Kind: l.Kind,
Billing: BillingMode(l.Kind), Status: l.Status, LastSeen: rfc3339(l.LastSeen),
Usage: u, CreatedAt: rfc3339(l.CreatedAt), UpdatedAt: rfc3339(l.UpdatedAt),
}
}
// devicesOf folds a user's links into the per-machine device projection, newest
// device first. A device's labels come from its most-recently-seen account.
func devicesOf(links []Link) []deviceView {
order := make([]string, 0)
byMachine := map[string]*deviceView{}
for _, l := range links {
d, ok := byMachine[l.Machine]
if !ok {
d = &deviceView{Machine: l.Machine, Host: l.Host, OS: l.OS, LastSeen: rfc3339(l.LastSeen)}
byMachine[l.Machine] = d
order = append(order, l.Machine)
}
// links arrive newest-first, so the first account seen carries the freshest
// device labels; keep them.
d.Accounts = append(d.Accounts, toLinkView(l))
}
out := make([]deviceView, 0, len(order))
for _, m := range order {
out = append(out, *byMachine[m])
}
return out
}
// ---- handlers ----
type registerReq struct {
Machine string `json:"machine"`
Host string `json:"host"`
OS string `json:"os"`
Provider string `json:"provider"`
Account string `json:"account"`
Plan string `json:"plan"`
Kind string `json:"kind"`
Usage json.RawMessage `json:"usage"`
}
func upsertLink(s *cloud.Service[state], c *zip.Ctx) error {
org, user, ok := caller(c)
if !ok {
return zip.ErrForbidden("X-Org-Id required")
}
var body registerReq
if err := c.Bind(&body); err != nil {
return err
}
machine := trim(body.Machine)
if machine == "" {
return zip.ErrBadRequest("machine is required")
}
provider := trim(body.Provider)
if provider == "" {
return zip.ErrBadRequest("provider is required")
}
if len(machine) > maxMachine || len(provider) > maxProvider ||
len(trim(body.Host)) > maxHost || len(trim(body.OS)) > maxOS ||
len(trim(body.Account)) > maxAccount || len(trim(body.Plan)) > maxPlan {
return zip.ErrBadRequest("field too long")
}
kind := trim(body.Kind)
if kind == "" {
kind = KindSubscription
}
if !validKind(kind) {
return zip.ErrBadRequest("kind must be subscription or apikey")
}
usageJSON, err := normalizeUsage(body.Usage)
if err != nil {
return err
}
id, err := genID("link")
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "rng: %v", err)
}
now := time.Now().Unix()
l := Link{
ID: id, Org: org, User: user, Machine: machine, Host: trim(body.Host), OS: trim(body.OS),
Provider: provider, Account: trim(body.Account), Plan: trim(body.Plan),
Kind: kind, Status: StatusLinked, LastSeen: now, Usage: usageJSON,
CreatedAt: now, UpdatedAt: now,
}
stored, err := s.State.store.Upsert(c.Context(), l)
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "persist: %v", err)
}
return c.JSON(http.StatusCreated, toLinkView(stored))
}
// normalizeUsage parses, clamps, and re-marshals a usage snapshot so the stored
// blob is bounded, well-formed, and carries only known fields. An absent usage is
// "" (a heartbeat that keeps the last good snapshot). A malformed or oversized
// usage is a 400 — the collector owns a valid projection.
func normalizeUsage(raw json.RawMessage) (string, error) {
if len(raw) == 0 || string(raw) == "null" {
return "", nil
}
if len(raw) > maxUsage {
return "", zip.ErrBadRequest("usage too large")
}
var u Usage
if err := json.Unmarshal(raw, &u); err != nil {
return "", zip.ErrBadRequest("usage must be a valid usage snapshot")
}
u.SessionPct = clampPct(u.SessionPct)
u.WeeklyPct = clampPct(u.WeeklyPct)
b, err := json.Marshal(u)
if err != nil {
return "", zip.ErrBadRequest("usage must be a valid usage snapshot")
}
return string(b), nil
}
func listLinks(s *cloud.Service[state], c *zip.Ctx) error {
org, user, ok := caller(c)
if !ok {
return zip.ErrForbidden("X-Org-Id required")
}
links, err := s.State.store.List(c.Context(), org, user)
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "list: %v", err)
}
views := make([]linkView, 0, len(links))
for _, l := range links {
views = append(views, toLinkView(l))
}
return c.JSON(http.StatusOK, map[string]any{"links": views, "devices": devicesOf(links)})
}
func getLink(s *cloud.Service[state], c *zip.Ctx) error {
org, user, ok := caller(c)
if !ok {
return zip.ErrForbidden("X-Org-Id required")
}
id := trim(c.Param("id"))
l, err := s.State.store.Get(c.Context(), org, user, id)
if err == errNotFound {
return zip.ErrNotFound("link not found")
}
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "get: %v", err)
}
return c.JSON(http.StatusOK, toLinkView(l))
}
func deviceDetail(s *cloud.Service[state], c *zip.Ctx) error {
org, user, ok := caller(c)
if !ok {
return zip.ErrForbidden("X-Org-Id required")
}
machine := trim(c.Param("machine"))
accounts, err := s.State.store.ListDevice(c.Context(), org, user, machine)
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "device: %v", err)
}
if len(accounts) == 0 {
return zip.ErrNotFound("device not found")
}
d := deviceView{Machine: machine, Host: accounts[0].Host, OS: accounts[0].OS, LastSeen: rfc3339(accounts[0].LastSeen)}
for _, a := range accounts {
d.Accounts = append(d.Accounts, toLinkView(a))
}
d.ActiveSessions = countActive(s, c.Context(), org, SessionMatch{Host: d.Host})
return c.JSON(http.StatusOK, d)
}
func routePlan(s *cloud.Service[state], c *zip.Ctx) error {
org, user, ok := caller(c)
if !ok {
return zip.ErrForbidden("X-Org-Id required")
}
linked, err := s.State.store.ListLinked(c.Context(), org, user)
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "route: %v", err)
}
return c.JSON(http.StatusOK, Plan(linked, time.Now()))
}
type revokeResp struct {
Revoked int `json:"revoked"`
SessionsStopped int `json:"sessionsStopped"`
Links []linkView `json:"links,omitempty"`
}
func revokeLink(s *cloud.Service[state], c *zip.Ctx) error {
org, user, ok := caller(c)
if !ok {
return zip.ErrForbidden("X-Org-Id required")
}
id := trim(c.Param("id"))
l, found, err := s.State.store.Revoke(c.Context(), org, user, id, time.Now().Unix())
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "revoke: %v", err)
}
if !found {
return zip.ErrNotFound("link not found")
}
stopped := stopSessions(s, c.Context(), org, SessionMatch{Host: l.Host, Provider: l.Provider, Account: l.Account})
return c.JSON(http.StatusOK, revokeResp{Revoked: 1, SessionsStopped: stopped, Links: []linkView{toLinkView(l)}})
}
func revokeDevice(s *cloud.Service[state], c *zip.Ctx) error {
org, user, ok := caller(c)
if !ok {
return zip.ErrForbidden("X-Org-Id required")
}
machine := trim(c.Param("machine"))
revoked, err := s.State.store.RevokeDevice(c.Context(), org, user, machine, time.Now().Unix())
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "revoke device: %v", err)
}
if len(revoked) == 0 {
return zip.ErrNotFound("device not found or already revoked")
}
// Stop every session on the device (all accounts).
stopped := stopSessions(s, c.Context(), org, SessionMatch{Host: revoked[0].Host})
views := make([]linkView, 0, len(revoked))
for _, l := range revoked {
views = append(views, toLinkView(l))
}
return c.JSON(http.StatusOK, revokeResp{Revoked: len(revoked), SessionsStopped: stopped, Links: views})
}
// stopSessions forwards to the sessions seam, tolerating a nil seam (unit test /
// no agents) and a seam error (a stop failure must not fail the revoke — the row
// is already revoked, which is the durable truth). Returns how many stopped.
func stopSessions(s *cloud.Service[state], ctx context.Context, org string, m SessionMatch) int {
if s.State.sessions == nil {
return 0
}
n, err := s.State.sessions.Stop(ctx, org, m)
if err != nil {
s.Log.Warn("link: stop sessions failed", "org", org, "err", err)
return 0
}
return n
}
func countActive(s *cloud.Service[state], ctx context.Context, org string, m SessionMatch) int {
if s.State.sessions == nil {
return 0
}
n, err := s.State.sessions.CountActive(ctx, org, m)
if err != nil {
return 0
}
return n
}
func genID(prefix string) (string, error) {
var b [16]byte
if _, err := rand.Read(b[:]); err != nil {
return "", err
}
return prefix + "_" + hex.EncodeToString(b[:]), nil
}
+314
View File
@@ -0,0 +1,314 @@
package link
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"testing"
"github.com/hanzoai/cloud"
"github.com/hanzoai/cloud/clients/agents"
luxlog "github.com/luxfi/log"
"github.com/zap-proto/zip"
)
func mountLink(t *testing.T) *zip.App {
t.Helper()
app := zip.New(zip.Config{Logger: luxlog.New("test")})
if err := Mount(app, cloud.Deps{Logger: luxlog.New("test"), DataDir: t.TempDir()}); err != nil {
t.Fatalf("Mount: %v", err)
}
t.Cleanup(func() { _ = Shutdown(context.Background()) })
return app
}
// req drives a request with a principal: X-Org-Id (the tenant) + X-User-Id (the
// validated subject, set only from a verified credential by SanitizeIdentity in
// prod — injected here as the gateway would). An empty user is the anonymous forge
// (org header, no validated principal) that every route must refuse.
func req(t *testing.T, app *zip.App, method, path, org, user string, body any) (int, []byte) {
t.Helper()
var r io.Reader
if body != nil {
b, _ := json.Marshal(body)
r = bytes.NewReader(b)
}
rq := httptest.NewRequest(method, path, r)
if body != nil {
rq.Header.Set("Content-Type", "application/json")
}
if org != "" {
rq.Header.Set("X-Org-Id", org)
}
if user != "" {
rq.Header.Set("X-User-Id", user)
}
resp, err := app.Fiber().Test(rq)
if err != nil {
t.Fatalf("Test %s %s: %v", method, path, err)
}
defer func() { _ = resp.Body.Close() }()
b, _ := io.ReadAll(resp.Body)
return resp.StatusCode, b
}
// TestFailClosedNoPrincipal: an org header with NO validated user (the off-gateway
// forge) is refused on every route — read AND write.
func TestFailClosedNoPrincipal(t *testing.T) {
app := mountLink(t)
for _, tc := range []struct {
method, path string
body any
}{
{http.MethodGet, "/v1/links", nil},
{http.MethodGet, "/v1/links/route", nil},
{http.MethodGet, "/v1/links/devices/m1", nil},
{http.MethodPost, "/v1/links", map[string]any{"machine": "m1", "provider": "claude"}},
{http.MethodDelete, "/v1/links/link_x", nil},
} {
if code, _ := req(t, app, tc.method, tc.path, "acme", "", tc.body); code != http.StatusForbidden {
t.Fatalf("%s %s with no validated principal want 403, got %d", tc.method, tc.path, code)
}
}
}
// TestRegisterListGetRevoke drives the full happy path AND asserts the
// subscription-vs-api-key billing distinction over the wire.
func TestRegisterListGetRevoke(t *testing.T) {
app := mountLink(t)
// Register a subscription account (Claude Max) with a usage snapshot.
code, body := req(t, app, http.MethodPost, "/v1/links", "acme", "alice", map[string]any{
"machine": "m1", "host": "box", "os": "darwin",
"provider": "claude", "account": "alice@x", "plan": "Claude Max", "kind": "subscription",
"usage": map[string]any{"sessionPct": 42, "weeklyPct": 12, "tokens": 1000},
})
if code != http.StatusCreated {
t.Fatalf("register want 201, got %d (%s)", code, body)
}
var sub linkView
_ = json.Unmarshal(body, &sub)
if sub.ID == "" || sub.Provider != "claude" || sub.Kind != "subscription" {
t.Fatalf("register shape: %+v", sub)
}
// A subscription bills the PLAN — never commerce.
if sub.Billing != BillingPlan {
t.Fatalf("a subscription account must carry billing=plan, got %q", sub.Billing)
}
// Register an api-key account (bills via commerce).
code, body = req(t, app, http.MethodPost, "/v1/links", "acme", "alice", map[string]any{
"machine": "m1", "host": "box", "provider": "hanzo", "account": "hk-1", "kind": "apikey",
})
if code != http.StatusCreated {
t.Fatalf("register apikey want 201, got %d (%s)", code, body)
}
var key linkView
_ = json.Unmarshal(body, &key)
if key.Billing != BillingCommerce {
t.Fatalf("an api-key account must carry billing=commerce, got %q", key.Billing)
}
// List returns both + a device projection grouping them under m1.
code, body = req(t, app, http.MethodGet, "/v1/links", "acme", "alice", nil)
var list struct {
Links []linkView `json:"links"`
Devices []deviceView `json:"devices"`
}
_ = json.Unmarshal(body, &list)
if code != http.StatusOK || len(list.Links) != 2 {
t.Fatalf("list want 2 links, got %d (%s)", len(list.Links), body)
}
if len(list.Devices) != 1 || list.Devices[0].Machine != "m1" || len(list.Devices[0].Accounts) != 2 {
t.Fatalf("device projection want m1 with 2 accounts, got %+v", list.Devices)
}
// The route plan puts the subscription first and carries the billing mode.
code, body = req(t, app, http.MethodGet, "/v1/links/route", "acme", "alice", nil)
var plan RoutePlan
_ = json.Unmarshal(body, &plan)
if code != http.StatusOK || len(plan.Candidates) != 2 {
t.Fatalf("route want 2 candidates, got %d (%s)", len(plan.Candidates), body)
}
if plan.Candidates[0].Kind != KindSubscription || plan.Candidates[0].Billing != BillingPlan {
t.Fatalf("route must prefer the subscription (plan billing), got %+v", plan.Candidates[0])
}
// Revoke the subscription → it drops out of the route; the api key remains.
code, body = req(t, app, http.MethodDelete, "/v1/links/"+sub.ID, "acme", "alice", nil)
var rev revokeResp
_ = json.Unmarshal(body, &rev)
if code != http.StatusOK || rev.Revoked != 1 {
t.Fatalf("revoke want 200 revoked=1, got %d %+v", code, rev)
}
code, body = req(t, app, http.MethodGet, "/v1/links/route", "acme", "alice", nil)
_ = json.Unmarshal(body, &plan)
if len(plan.Candidates) != 1 || plan.Candidates[0].Kind != KindAPIKey {
t.Fatalf("after revoke only the api key routes, got %+v", plan.Candidates)
}
// The revoked link is still listed (retained) but marked revoked.
code, body = req(t, app, http.MethodGet, "/v1/links/"+sub.ID, "acme", "alice", nil)
var got linkView
_ = json.Unmarshal(body, &got)
if code != http.StatusOK || got.Status != StatusRevoked {
t.Fatalf("revoked link must remain gettable as revoked, got %d %q", code, got.Status)
}
}
// TestHTTPUserAndOrgIsolation: a second user in the same org, AND a second org,
// see/get/revoke none of the first user's links.
func TestHTTPUserAndOrgIsolation(t *testing.T) {
app := mountLink(t)
code, body := req(t, app, http.MethodPost, "/v1/links", "acme", "alice", map[string]any{
"machine": "m1", "provider": "claude", "account": "a", "kind": "subscription",
})
if code != http.StatusCreated {
t.Fatalf("alice register want 201, got %d", code)
}
var al linkView
_ = json.Unmarshal(body, &al)
// bob (same org) sees nothing and cannot get/revoke alice's link by id.
_, body = req(t, app, http.MethodGet, "/v1/links", "acme", "bob", nil)
var bobList struct {
Links []linkView `json:"links"`
}
_ = json.Unmarshal(body, &bobList)
if len(bobList.Links) != 0 {
t.Fatalf("bob must see zero links, got %d", len(bobList.Links))
}
if code, _ := req(t, app, http.MethodGet, "/v1/links/"+al.ID, "acme", "bob", nil); code != http.StatusNotFound {
t.Fatalf("bob GET alice's link want 404, got %d", code)
}
if code, _ := req(t, app, http.MethodDelete, "/v1/links/"+al.ID, "acme", "bob", nil); code != http.StatusNotFound {
t.Fatalf("bob DELETE alice's link want 404, got %d", code)
}
// evil org sees nothing either.
_, body = req(t, app, http.MethodGet, "/v1/links", "evil", "alice", nil)
var evilList struct {
Links []linkView `json:"links"`
}
_ = json.Unmarshal(body, &evilList)
if len(evilList.Links) != 0 {
t.Fatalf("evil org must see zero of acme's links, got %d", len(evilList.Links))
}
// alice's link is intact after the foreign attempts.
if code, _ := req(t, app, http.MethodGet, "/v1/links/"+al.ID, "acme", "alice", nil); code != http.StatusOK {
t.Fatalf("alice's own link must survive, got %d", code)
}
}
// TestHTTPInputValidation: a register missing a required field, a bad kind, or an
// oversized usage is a clean 400.
func TestHTTPInputValidation(t *testing.T) {
app := mountLink(t)
bad := []map[string]any{
{"provider": "claude"}, // no machine
{"machine": "m1"}, // no provider
{"machine": "m1", "provider": "claude", "kind": "freeloader"}, // bad kind
}
for _, b := range bad {
if code, _ := req(t, app, http.MethodPost, "/v1/links", "acme", "alice", b); code != http.StatusBadRequest {
t.Fatalf("invalid register %+v want 400, got %d", b, code)
}
}
// A device with no accounts is a 404.
if code, _ := req(t, app, http.MethodGet, "/v1/links/devices/nope", "acme", "alice", nil); code != http.StatusNotFound {
t.Fatalf("unknown device want 404, got %d", code)
}
}
// TestRevokeStopsSessions is the end-to-end proof of the login-out contract: with
// BOTH the agents session plane and the link registry mounted, revoking a link
// stops the live sessions that ran under that account — via the REAL adapter, not
// a fake. It also proves the session↔account tag (a session carries its
// provider/account) and org isolation of the stop (another org's identical session
// is untouched).
func TestRevokeStopsSessions(t *testing.T) {
dir := t.TempDir()
app := zip.New(zip.Config{Logger: luxlog.New("test")})
deps := cloud.Deps{Logger: luxlog.New("test"), DataDir: dir}
if err := agents.Mount(app, deps); err != nil {
t.Fatalf("agents.Mount: %v", err)
}
t.Cleanup(func() { _ = agents.Shutdown(context.Background()) })
if err := Mount(app, deps); err != nil {
t.Fatalf("link.Mount: %v", err)
}
t.Cleanup(func() { _ = Shutdown(context.Background()) })
// A live session tagged with the account it runs under, on host "box1".
code, body := req(t, app, http.MethodPost, "/v1/agents/sessions", "acme", "alice", map[string]any{
"agent": "dev", "host": "box1", "provider": "claude", "account": "alice@x",
})
if code != http.StatusCreated {
t.Fatalf("register session want 201, got %d (%s)", code, body)
}
var sess struct {
ID string `json:"id"`
Status string `json:"status"`
Provider string `json:"provider"`
Account string `json:"account"`
}
_ = json.Unmarshal(body, &sess)
if sess.Provider != "claude" || sess.Account != "alice@x" {
t.Fatalf("session must carry its account tag, got %+v", sess)
}
// The SAME account in ANOTHER org — must be untouched by acme's revoke.
code, body = req(t, app, http.MethodPost, "/v1/agents/sessions", "evil", "mallory", map[string]any{
"agent": "dev", "host": "box1", "provider": "claude", "account": "alice@x",
})
if code != http.StatusCreated {
t.Fatalf("evil session want 201, got %d", code)
}
var evil struct {
ID string `json:"id"`
}
_ = json.Unmarshal(body, &evil)
// Register the link for that account, then revoke it (log out).
code, body = req(t, app, http.MethodPost, "/v1/links", "acme", "alice", map[string]any{
"machine": "m1", "host": "box1", "provider": "claude", "account": "alice@x", "kind": "subscription",
})
if code != http.StatusCreated {
t.Fatalf("register link want 201, got %d", code)
}
var l linkView
_ = json.Unmarshal(body, &l)
code, body = req(t, app, http.MethodDelete, "/v1/links/"+l.ID, "acme", "alice", nil)
var rev revokeResp
_ = json.Unmarshal(body, &rev)
if code != http.StatusOK {
t.Fatalf("revoke want 200, got %d (%s)", code, body)
}
if rev.SessionsStopped != 1 {
t.Fatalf("revoke must stop the 1 session under that account, got %d", rev.SessionsStopped)
}
// The acme session is now terminal (stopped).
_, body = req(t, app, http.MethodGet, "/v1/agents/sessions/"+sess.ID, "acme", "alice", nil)
var after struct {
Status string `json:"status"`
}
_ = json.Unmarshal(body, &after)
if after.Status != "error" {
t.Fatalf("the revoked account's session must be stopped (terminal), got %q", after.Status)
}
// The evil org's identical session is UNTOUCHED (org-scoped stop).
_, body = req(t, app, http.MethodGet, "/v1/agents/sessions/"+evil.ID, "evil", "mallory", nil)
var evilAfter struct {
Status string `json:"status"`
}
_ = json.Unmarshal(body, &evilAfter)
if evilAfter.Status != "running" {
t.Fatalf("another org's session must survive acme's revoke, got %q", evilAfter.Status)
}
}
+163
View File
@@ -0,0 +1,163 @@
// Package link is the unified AI login manager's registry: the org+user-scoped
// record of WHICH provider accounts (Claude Max, ChatGPT Plus, a Hanzo API key,
// a raw provider key) a developer has signed into, ON WHICH MACHINES, with each
// account's latest usage snapshot. It is the cross-machine view console renders
// as "AI Providers / Accounts" and the source the redundancy route policy reads.
//
// THE ATOM is a Link: one (user, device, provider, account) binding plus its
// kind, status, last-seen, and the latest usage snapshot. A device is the
// (machine, host, os) projection shared by a machine's Links — not a separate
// stored entity, so there is no device/link join and no orphan-device GC. A
// user's accounts across every machine are just that user's Links.
//
// NO SECRET LIVES HERE. The provider's OAuth token / API key stays device-local
// (exactly as @hanzo/usage keeps it, reading each provider's own login to meter
// usage); the registry holds LINK METADATA + usage snapshots only. The collector
// (@hanzo/usage's reporter) registers a Link and pushes its usage using the SAME
// IAM bearer it already carries — never a provider secret.
//
// BILLING FOLLOWS THE ACCOUNT, not the registry. This store never touches
// commerce — it holds no metering client, so it is structurally incapable of
// creating a charge. A subscription account's inference bills the user's monthly
// plan (metered here for visibility only); an api-key / hanzo account's inference
// bills via commerce on the existing gateway path, unchanged. Each Link carries
// the mode that says which (BillingMode), so a usage event's billing is explicit.
//
// ISOLATION: org is the tenant key and user is the owner key; every read and
// write leads with `org=? AND user=?` bound predicates, so a caller sees and
// mutates only their OWN accounts within their OWN org — fail-closed.
package link
import "strings"
// Link kinds — how an account is credentialed, which decides how its usage bills.
const (
// KindSubscription is a provider account signed in with the user's own
// subscription login (Claude Max, ChatGPT Plus). Its inference bills the
// user's monthly plan; the registry meters it for visibility only.
KindSubscription = "subscription"
// KindAPIKey is an account credentialed by an API key (a raw provider key,
// or a Hanzo hk- key). Its inference bills via commerce on the gateway path.
KindAPIKey = "apikey"
)
// Link statuses. linked is active; revoked is a logged-out account whose sessions
// were stopped. A revoked Link is retained (not deleted) so its usage history and
// the audit trail survive a log-out.
const (
StatusLinked = "linked"
StatusRevoked = "revoked"
)
// Billing modes — the value a Link and a route candidate carry so "how does this
// account's usage bill" is explicit, never inferred at the charge site.
const (
// BillingPlan: the user's own subscription pays; NO commerce charge.
BillingPlan = "plan"
// BillingCommerce: usage bills via commerce (the gateway meter), as today.
BillingCommerce = "commerce"
)
// Field bounds. Every string is a short identifier or label, never a document;
// the caps keep a hostile client from bloating the per-org SQLite and match the
// session plane's discipline.
const (
maxUser = 256
maxMachine = 256
maxHost = 256
maxOS = 64
maxProvider = 64
maxAccount = 256
maxPlan = 128
maxUsage = 8 * 1024 // a usage snapshot is a handful of numbers, not a blob
)
// Link is one (user, device, provider, account) binding. Tenant isolation is the
// (org, user) pair, enforced on every query. It never stores a secret.
type Link struct {
ID string
Org string
User string // the owning subject (validated principal); a user sees only their own
Machine string // stable machine id (the device key)
Host string // hostname (device label)
OS string // platform label (darwin|linux|windows|…)
Provider string // matches @hanzo/usage providerRegistry id (claude|codex|hanzo|openai|…)
Account string // the subscription/account id or label (e.g. an account email)
Plan string // the plan name from the provider identity (e.g. "Claude Max"); display only
Kind string // subscription | apikey
Status string // linked | revoked
LastSeen int64 // unix; bumped on every usage report
Usage string // JSON of the latest Usage projection ("" until first report)
CreatedAt int64
UpdatedAt int64
}
// Usage is the projection of a provider's @hanzo/usage UsageSnapshot the collector
// pushes and the dashboard renders: the rate-limit windows, token totals, and
// spend. It is stored as JSON on the Link and parsed for the route policy's
// headroom. Money is USD cents end-to-end. Spend is always 0 for a pure
// subscription account (there is no per-call charge — the plan is flat).
type Usage struct {
SessionPct float64 `json:"sessionPct"` // primary window used %, 0..100
WeeklyPct float64 `json:"weeklyPct"` // secondary window used %, 0..100
ResetsAt string `json:"resetsAt,omitempty"` // primary window reset, RFC3339
Tokens int64 `json:"tokens"` // absolute token total when known
InputTokens int64 `json:"inputTokens,omitempty"` //
OutputTokens int64 `json:"outputTokens,omitempty"` //
SpendCents int64 `json:"spendCents"` // provider spend (0 for a subscription)
Currency string `json:"currency,omitempty"` //
Confidence string `json:"confidence,omitempty"` // exact|estimated|percentOnly|unknown
UpdatedAt string `json:"updatedAt,omitempty"` // snapshot time, RFC3339
}
// validKind reports whether k is a known link kind.
func validKind(k string) bool { return k == KindSubscription || k == KindAPIKey }
// BillingMode returns how an account of this kind bills its usage: a subscription
// bills the user's plan (no commerce charge), everything else bills via commerce.
// It is the ONE place the subscription-vs-api-key distinction is decided, so a
// usage event's billing is a pure function of the account, never re-derived.
func BillingMode(kind string) string {
if kind == KindSubscription {
return BillingPlan
}
return BillingCommerce
}
// headroomPct is the remaining capacity of the tighter of a snapshot's two
// rate-limit windows (100 - the higher used%), clamped to [0,100]. A Link with no
// usage snapshot is treated as fully available (100) — absence of data is not
// evidence of exhaustion. This is the value the route policy orders candidates by.
func headroomPct(u *Usage) float64 {
if u == nil {
return 100
}
used := u.SessionPct
if u.WeeklyPct > used {
used = u.WeeklyPct
}
h := 100 - used
if h < 0 {
return 0
}
if h > 100 {
return 100
}
return h
}
// clampPct bounds a percent to [0,100]; a provider that reports a nonsense value
// can never poison the headroom ordering.
func clampPct(v float64) float64 {
if v < 0 {
return 0
}
if v > 100 {
return 100
}
return v
}
// trim is strings.TrimSpace, aliased for terse boundary code.
func trim(s string) string { return strings.TrimSpace(s) }
+91
View File
@@ -0,0 +1,91 @@
package link
import (
"encoding/json"
"strings"
"testing"
)
// TestBillingModeIsTheSubscriptionDistinction: the ONE decision that separates a
// subscription (bills the user's plan — NO commerce charge) from an api-key
// account (bills via commerce). An unknown kind is never silently "free".
func TestBillingModeIsTheSubscriptionDistinction(t *testing.T) {
if BillingMode(KindSubscription) != BillingPlan {
t.Fatalf("a subscription must bill the plan, not commerce")
}
if BillingMode(KindAPIKey) != BillingCommerce {
t.Fatalf("an api-key account must bill via commerce")
}
if BillingMode("weird") != BillingCommerce {
t.Fatalf("an unknown kind must default to a real charge, never free")
}
}
func TestHeadroom(t *testing.T) {
if headroomPct(nil) != 100 {
t.Fatalf("no snapshot = full headroom")
}
if got := headroomPct(&Usage{SessionPct: 30, WeeklyPct: 70}); got != 30 {
t.Fatalf("headroom is the tighter window (100-70=30), got %v", got)
}
if got := headroomPct(&Usage{SessionPct: 100}); got != 0 {
t.Fatalf("an exhausted window = 0 headroom, got %v", got)
}
if got := headroomPct(&Usage{SessionPct: 150}); got != 0 {
t.Fatalf("over-100 used clamps headroom to 0, got %v", got)
}
if got := headroomPct(&Usage{SessionPct: -20}); got != 100 {
t.Fatalf("a nonsense negative used clamps to full headroom, got %v", got)
}
}
// TestNormalizeUsage: absent → "" (heartbeat), valid → clamped + re-marshalled,
// malformed / oversized → an error (the collector owns a valid projection).
func TestNormalizeUsage(t *testing.T) {
if out, err := normalizeUsage(nil); err != nil || out != "" {
t.Fatalf("absent usage → empty, got %q %v", out, err)
}
if out, err := normalizeUsage(json.RawMessage("null")); err != nil || out != "" {
t.Fatalf("null usage → empty, got %q %v", out, err)
}
out, err := normalizeUsage(json.RawMessage(`{"sessionPct":150,"weeklyPct":-5,"tokens":42}`))
if err != nil {
t.Fatalf("valid usage: %v", err)
}
var u Usage
_ = json.Unmarshal([]byte(out), &u)
if u.SessionPct != 100 || u.WeeklyPct != 0 || u.Tokens != 42 {
t.Fatalf("usage must be clamped [0,100] and preserved, got %+v", u)
}
if _, err := normalizeUsage(json.RawMessage(`{not json`)); err == nil {
t.Fatalf("malformed usage must error")
}
big := json.RawMessage(`{"resetsAt":"` + strings.Repeat("x", maxUsage) + `"}`)
if _, err := normalizeUsage(big); err == nil {
t.Fatalf("oversized usage must error")
}
}
// TestDevicesOf groups a user's links by machine (device projection), newest
// device first, carrying each machine's accounts.
func TestDevicesOf(t *testing.T) {
links := []Link{
{ID: "1", Machine: "m1", Host: "h1", Provider: "claude", Kind: KindSubscription, LastSeen: 200},
{ID: "2", Machine: "m1", Host: "h1", Provider: "codex", Kind: KindSubscription, LastSeen: 190},
{ID: "3", Machine: "m2", Host: "h2", Provider: "hanzo", Kind: KindAPIKey, LastSeen: 180},
}
ds := devicesOf(links)
if len(ds) != 2 {
t.Fatalf("want 2 devices, got %d", len(ds))
}
if ds[0].Machine != "m1" || len(ds[0].Accounts) != 2 {
t.Fatalf("m1 groups its 2 accounts, got %+v", ds[0])
}
if ds[1].Machine != "m2" || len(ds[1].Accounts) != 1 {
t.Fatalf("m2 has 1 account, got %+v", ds[1])
}
// The account views carry the derived billing mode.
if ds[0].Accounts[0].Billing != BillingPlan || ds[1].Accounts[0].Billing != BillingCommerce {
t.Fatalf("device views must carry each account's billing mode")
}
}
+133
View File
@@ -0,0 +1,133 @@
package link
import (
"encoding/json"
"sort"
"time"
)
// route.go is the redundancy seam: a PURE policy that turns a user's linked
// accounts into an ordered list of routing candidates, so a caller (the ai
// gateway / enso router) can fail over across accounts — two Claude Max
// subscriptions for redundancy, then the metered API as the always-available
// backstop — and know how each candidate BILLS before it dials.
//
// This is the policy + its types (the SEAM). EXECUTION — actually dialing a
// provider, detecting a live 429, and advancing to the next candidate — is the
// gateway's job (the deferred failover-execution increment). The policy reads the
// registry's own usage snapshots (the rate-limit headroom @hanzo/usage already
// meters), never a live provider probe, so it is a total function of the Links.
// RouteCandidate is one account the policy would route to, in preference order,
// annotated with how it bills and whether it currently has rate-limit headroom.
type RouteCandidate struct {
Provider string `json:"provider"`
Account string `json:"account,omitempty"`
Plan string `json:"plan,omitempty"`
Kind string `json:"kind"` // subscription | apikey
Billing string `json:"billing"` // plan | commerce (BillingMode(Kind))
Available bool `json:"available"`
HeadroomPct float64 `json:"headroomPct"` // remaining capacity 0..100
Machine string `json:"machine,omitempty"`
Host string `json:"host,omitempty"`
LinkID string `json:"linkId"`
Reason string `json:"reason,omitempty"` // why unavailable, when Available=false
}
// RoutePlan is the ordered redundancy plan for a user's accounts: the candidates
// in preference order plus the primary the caller would try first.
type RoutePlan struct {
Candidates []RouteCandidate `json:"candidates"`
Primary *RouteCandidate `json:"primary,omitempty"`
GeneratedAt string `json:"generatedAt"`
}
// parseUsage decodes a Link's stored usage JSON, or nil when there is none / it
// is malformed (treated as "no snapshot" → full headroom, never a crash).
func parseUsage(raw string) *Usage {
if raw == "" {
return nil
}
var u Usage
if err := json.Unmarshal([]byte(raw), &u); err != nil {
return nil
}
return &u
}
// candidateOf projects a Link into a RouteCandidate: headroom from its usage
// snapshot, availability from that headroom (a fully-consumed window → not
// routable right now), and the billing mode from its kind.
func candidateOf(l Link) RouteCandidate {
h := headroomPct(parseUsage(l.Usage))
c := RouteCandidate{
Provider: l.Provider, Account: l.Account, Plan: l.Plan, Kind: l.Kind,
Billing: BillingMode(l.Kind), HeadroomPct: h,
Machine: l.Machine, Host: l.Host, LinkID: l.ID,
Available: h > 0,
}
if !c.Available {
c.Reason = "rate limit reached"
}
return c
}
// Plan builds the redundancy route plan from a user's LINKED accounts. Ordering:
// subscription accounts first (the flat-rate pool the user already pays for —
// used first, and failed over between for redundancy: two Claude Max accounts),
// then api-key/hanzo accounts (the metered API route). Within each group the
// order is (available first, then most headroom, then most-recently seen). The
// primary is the first available candidate; if every account is rate-limited it
// falls back to the first api-key account (the pay-per-call backstop is always
// usable) — else nil, an honest "nothing routable right now".
//
// Only linked accounts should be passed (Store.ListLinked); a revoked account is
// never a candidate. The function is pure and total: any Links in, one plan out.
func Plan(links []Link, now time.Time) RoutePlan {
subs := make([]RouteCandidate, 0, len(links))
keys := make([]RouteCandidate, 0, len(links))
for _, l := range links {
if l.Status != StatusLinked {
continue
}
c := candidateOf(l)
if l.Kind == KindSubscription {
subs = append(subs, c)
} else {
keys = append(keys, c)
}
}
byPreference := func(cs []RouteCandidate) {
sort.SliceStable(cs, func(i, j int) bool {
if cs[i].Available != cs[j].Available {
return cs[i].Available // available first
}
return cs[i].HeadroomPct > cs[j].HeadroomPct // then most headroom
})
}
byPreference(subs)
byPreference(keys)
candidates := append(append(make([]RouteCandidate, 0, len(subs)+len(keys)), subs...), keys...)
plan := RoutePlan{Candidates: candidates, GeneratedAt: now.UTC().Format(time.RFC3339)}
plan.Primary = pickPrimary(candidates)
return plan
}
// pickPrimary returns the first available candidate; failing that, the first
// api-key candidate (the metered backstop is usable even when the subscriptions
// are exhausted); failing that, nil.
func pickPrimary(cs []RouteCandidate) *RouteCandidate {
for i := range cs {
if cs[i].Available {
return &cs[i]
}
}
for i := range cs {
if cs[i].Kind == KindAPIKey {
return &cs[i]
}
}
return nil
}
+86
View File
@@ -0,0 +1,86 @@
package link
import (
"encoding/json"
"testing"
"time"
)
func lk(provider, account, kind string, sessionPct float64) Link {
u, _ := json.Marshal(Usage{SessionPct: sessionPct})
return Link{
ID: "id-" + account, Provider: provider, Account: account, Kind: kind,
Status: StatusLinked, Usage: string(u),
}
}
// TestRoutePlanRedundancyAndBilling: subscriptions (the flat-rate redundancy pool)
// come first, ordered by headroom; the api key is the trailing metered route; each
// candidate carries its billing mode; the primary is the most-available account.
func TestRoutePlanRedundancyAndBilling(t *testing.T) {
links := []Link{
lk("claude", "maxA", KindSubscription, 80), // headroom 20
lk("claude", "maxB", KindSubscription, 10), // headroom 90 — two Claude Max, redundancy
lk("hanzo", "hk", KindAPIKey, 0), // headroom 100 — the API route
}
p := Plan(links, time.Unix(0, 0))
if len(p.Candidates) != 3 {
t.Fatalf("want 3 candidates, got %d", len(p.Candidates))
}
if p.Candidates[0].Account != "maxB" || p.Candidates[1].Account != "maxA" {
t.Fatalf("subscriptions must order by headroom desc, got %+v", p.Candidates)
}
if p.Candidates[2].Provider != "hanzo" {
t.Fatalf("the api key must come after the subscriptions, got %+v", p.Candidates[2])
}
if p.Candidates[0].Billing != BillingPlan || p.Candidates[2].Billing != BillingCommerce {
t.Fatalf("each candidate must carry its billing mode")
}
if p.Primary == nil || p.Primary.Account != "maxB" {
t.Fatalf("primary should be the most-available subscription (maxB), got %+v", p.Primary)
}
}
// TestRouteFailoverToAPIWhenSubsExhausted: when every subscription is rate-limited,
// the primary fails over to the pay-per-call api key (always usable); the exhausted
// subscriptions remain listed as unavailable candidates with a reason.
func TestRouteFailoverToAPIWhenSubsExhausted(t *testing.T) {
links := []Link{
lk("claude", "maxA", KindSubscription, 100),
lk("claude", "maxB", KindSubscription, 100),
lk("hanzo", "hk", KindAPIKey, 50),
}
p := Plan(links, time.Unix(0, 0))
if p.Primary == nil || p.Primary.Kind != KindAPIKey {
t.Fatalf("primary must fail over to the api key, got %+v", p.Primary)
}
for _, c := range p.Candidates {
if c.Kind == KindSubscription {
if c.Available {
t.Fatalf("an exhausted subscription must be unavailable, got %+v", c)
}
if c.Reason == "" {
t.Fatalf("an unavailable candidate must carry a reason")
}
}
}
}
// TestRouteAllExhaustedNoPrimary: all subscriptions exhausted and no api key →
// an honest nil primary (nothing routable right now), never a fabricated pick.
func TestRouteAllExhaustedNoPrimary(t *testing.T) {
p := Plan([]Link{lk("claude", "maxA", KindSubscription, 100)}, time.Unix(0, 0))
if p.Primary != nil {
t.Fatalf("no routable account → primary must be nil, got %+v", p.Primary)
}
}
// TestRouteExcludesRevoked: a revoked (logged-out) account is never a candidate.
func TestRouteExcludesRevoked(t *testing.T) {
l := lk("claude", "maxA", KindSubscription, 10)
l.Status = StatusRevoked
p := Plan([]Link{l}, time.Unix(0, 0))
if len(p.Candidates) != 0 {
t.Fatalf("a revoked account must never be a route candidate, got %+v", p.Candidates)
}
}
+251
View File
@@ -0,0 +1,251 @@
package link
import (
"context"
"database/sql"
"errors"
"fmt"
// github.com/hanzoai/sqlite is the ONE Hanzo SQLite driver (registers the
// "sqlite" name under both build tags: cgo→mattn+SQLCipher, !cgo→modernc).
"github.com/hanzoai/cloud/cek"
_ "github.com/hanzoai/sqlite"
)
var errNotFound = errors.New("link: not found")
// Store is the login-manager database. ONE SQLite file ({DataDir}/link.db) holds
// every org's Links; tenancy is the (org, subject) pair. It holds NO metering
// client — it is structurally incapable of charging commerce.
type Store struct {
db *sql.DB
}
func openStore(path string) (*Store, error) {
db, err := cek.Open(path)
if err != nil {
return nil, fmt.Errorf("open sqlite %q: %w", path, err)
}
db.SetMaxOpenConns(1)
for _, pragma := range []string{
"PRAGMA busy_timeout=5000",
"PRAGMA journal_mode=WAL",
"PRAGMA foreign_keys=ON",
} {
if _, err := db.Exec(pragma); err != nil {
_ = db.Close()
return nil, fmt.Errorf("pragma %q: %w", pragma, err)
}
}
s := &Store{db: db}
if err := s.migrate(); err != nil {
_ = db.Close()
return nil, err
}
return s, nil
}
func (s *Store) migrate() error {
const ddl = `
CREATE TABLE IF NOT EXISTS links (
id TEXT PRIMARY KEY,
org TEXT NOT NULL,
subject TEXT NOT NULL,
machine TEXT NOT NULL,
host TEXT NOT NULL DEFAULT '',
os TEXT NOT NULL DEFAULT '',
provider TEXT NOT NULL,
account TEXT NOT NULL DEFAULT '',
plan TEXT NOT NULL DEFAULT '',
kind TEXT NOT NULL DEFAULT 'subscription',
status TEXT NOT NULL DEFAULT 'linked',
last_seen INTEGER NOT NULL DEFAULT 0,
usage TEXT NOT NULL DEFAULT '',
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
-- One Link per (tenant, owner, device, provider, account): the identity a
-- collector upserts against, so a re-report updates in place instead of duping.
CREATE UNIQUE INDEX IF NOT EXISTS ux_links_identity
ON links(org, subject, machine, provider, account);
-- The list projection: a user's links newest-first within their org.
CREATE INDEX IF NOT EXISTS ix_links_owner ON links(org, subject, updated_at);
-- The device projection: a machine's accounts within a user's scope.
CREATE INDEX IF NOT EXISTS ix_links_device ON links(org, subject, machine);
`
if _, err := s.db.Exec(ddl); err != nil {
return fmt.Errorf("migrate: %w", err)
}
return nil
}
func (s *Store) Close() error { return s.db.Close() }
const linkCols = `id,org,subject,machine,host,os,provider,account,plan,kind,status,last_seen,usage,created_at,updated_at`
func scanLink(sc interface{ Scan(...any) error }) (Link, error) {
var x Link
err := sc.Scan(&x.ID, &x.Org, &x.User, &x.Machine, &x.Host, &x.OS, &x.Provider,
&x.Account, &x.Plan, &x.Kind, &x.Status, &x.LastSeen, &x.Usage, &x.CreatedAt, &x.UpdatedAt)
return x, err
}
// Upsert registers a Link, keyed by its (org, subject, machine, provider,
// account) identity. A first report INSERTs (using x.ID / x.CreatedAt); a repeat
// UPDATEs in place — bumping last_seen, refreshing the device labels/plan/kind,
// re-activating a revoked account (status→linked), and replacing the usage
// snapshot ONLY when a fresh one is supplied (an empty usage on a heartbeat keeps
// the last good snapshot, mirroring @hanzo/usage's keep-stale-over-flapping rule).
// It returns the stored row (with its authoritative id, which on a repeat is the
// original, not x.ID). Org+subject are part of the identity, so a caller can only
// ever write within their OWN (org, subject) scope.
func (s *Store) Upsert(ctx context.Context, x Link) (Link, error) {
_, err := s.db.ExecContext(ctx,
`INSERT INTO links (`+linkCols+`)
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
ON CONFLICT(org, subject, machine, provider, account) DO UPDATE SET
host=excluded.host,
os=excluded.os,
plan=excluded.plan,
kind=excluded.kind,
status='linked',
last_seen=excluded.last_seen,
usage=CASE WHEN excluded.usage<>'' THEN excluded.usage ELSE links.usage END,
updated_at=excluded.updated_at`,
x.ID, x.Org, x.User, x.Machine, x.Host, x.OS, x.Provider, x.Account, x.Plan,
x.Kind, x.Status, x.LastSeen, x.Usage, x.CreatedAt, x.UpdatedAt)
if err != nil {
return Link{}, fmt.Errorf("upsert link: %w", err)
}
// Read back by identity so the returned row carries the authoritative id
// (original on a repeat) and the merged usage. Single connection serializes
// this with the write above.
return s.getByIdentity(ctx, x.Org, x.User, x.Machine, x.Provider, x.Account)
}
func (s *Store) getByIdentity(ctx context.Context, org, subject, machine, provider, account string) (Link, error) {
row := s.db.QueryRowContext(ctx,
`SELECT `+linkCols+` FROM links
WHERE org=? AND subject=? AND machine=? AND provider=? AND account=?`,
org, subject, machine, provider, account)
x, err := scanLink(row)
if errors.Is(err, sql.ErrNoRows) {
return Link{}, errNotFound
}
if err != nil {
return Link{}, fmt.Errorf("get by identity: %w", err)
}
return x, nil
}
// Get returns one Link by id within the caller's (org, subject) scope, or
// errNotFound. The (org, subject, id) triple is the key, so another tenant's or
// another user's id resolves to errNotFound — never a cross-scope read.
func (s *Store) Get(ctx context.Context, org, subject, id string) (Link, error) {
row := s.db.QueryRowContext(ctx,
`SELECT `+linkCols+` FROM links WHERE org=? AND subject=? AND id=?`, org, subject, id)
x, err := scanLink(row)
if errors.Is(err, sql.ErrNoRows) {
return Link{}, errNotFound
}
if err != nil {
return Link{}, fmt.Errorf("get link: %w", err)
}
return x, nil
}
// List returns every Link the (org, subject) owns, most-recently-updated first
// (revoked rows included, so the dashboard can show recent log-outs).
func (s *Store) List(ctx context.Context, org, subject string) ([]Link, error) {
return s.query(ctx,
`SELECT `+linkCols+` FROM links WHERE org=? AND subject=? ORDER BY updated_at DESC, id ASC`,
org, subject)
}
// ListDevice returns the accounts on one machine within the caller's scope,
// most-recently-updated first.
func (s *Store) ListDevice(ctx context.Context, org, subject, machine string) ([]Link, error) {
return s.query(ctx,
`SELECT `+linkCols+` FROM links WHERE org=? AND subject=? AND machine=? ORDER BY updated_at DESC, id ASC`,
org, subject, machine)
}
// ListLinked returns only the ACTIVE (status=linked) accounts the (org, subject)
// owns — the set the route policy considers. Revoked accounts are excluded so a
// logged-out account is never a routing candidate.
func (s *Store) ListLinked(ctx context.Context, org, subject string) ([]Link, error) {
return s.query(ctx,
`SELECT `+linkCols+` FROM links WHERE org=? AND subject=? AND status=? ORDER BY updated_at DESC, id ASC`,
org, subject, StatusLinked)
}
func (s *Store) query(ctx context.Context, q string, args ...any) ([]Link, error) {
rows, err := s.db.QueryContext(ctx, q, args...)
if err != nil {
return nil, fmt.Errorf("query links: %w", err)
}
defer func() { _ = rows.Close() }()
var out []Link
for rows.Next() {
x, err := scanLink(rows)
if err != nil {
return nil, fmt.Errorf("scan link: %w", err)
}
out = append(out, x)
}
return out, rows.Err()
}
// Revoke marks one Link revoked within the caller's scope and returns it (so the
// handler can stop the sessions that ran under that account). ok=false when no
// such Link exists in this (org, subject) — a cross-scope id can neither revoke
// nor probe. An already-revoked Link is returned as-is (idempotent).
func (s *Store) Revoke(ctx context.Context, org, subject, id string, now int64) (Link, bool, error) {
x, err := s.Get(ctx, org, subject, id)
if errors.Is(err, errNotFound) {
return Link{}, false, nil
}
if err != nil {
return Link{}, false, err
}
if x.Status == StatusRevoked {
return x, true, nil
}
res, err := s.db.ExecContext(ctx,
`UPDATE links SET status=?, updated_at=? WHERE org=? AND subject=? AND id=?`,
StatusRevoked, now, org, subject, id)
if err != nil {
return Link{}, false, fmt.Errorf("revoke link: %w", err)
}
if n, _ := res.RowsAffected(); n == 0 {
return Link{}, false, nil
}
x.Status = StatusRevoked
x.UpdatedAt = now
return x, true, nil
}
// RevokeDevice marks EVERY still-linked account on one machine revoked within the
// caller's scope and returns the rows it revoked (so their sessions can be
// stopped). A device with no linked accounts revokes nothing (empty slice).
func (s *Store) RevokeDevice(ctx context.Context, org, subject, machine string, now int64) ([]Link, error) {
linked, err := s.query(ctx,
`SELECT `+linkCols+` FROM links WHERE org=? AND subject=? AND machine=? AND status=?`,
org, subject, machine, StatusLinked)
if err != nil {
return nil, err
}
if len(linked) == 0 {
return nil, nil
}
if _, err := s.db.ExecContext(ctx,
`UPDATE links SET status=?, updated_at=? WHERE org=? AND subject=? AND machine=? AND status=?`,
StatusRevoked, now, org, subject, machine, StatusLinked); err != nil {
return nil, fmt.Errorf("revoke device: %w", err)
}
for i := range linked {
linked[i].Status = StatusRevoked
linked[i].UpdatedAt = now
}
return linked, nil
}
+164
View File
@@ -0,0 +1,164 @@
package link
import (
"context"
"path/filepath"
"testing"
"time"
)
func testStore(t *testing.T) *Store {
t.Helper()
s, err := openStore(filepath.Join(t.TempDir(), "link.db"))
if err != nil {
t.Fatalf("openStore: %v", err)
}
t.Cleanup(func() { _ = s.Close() })
return s
}
func mkLink(org, user, machine, provider, account, kind string) Link {
now := time.Now().Unix()
id, _ := genID("link")
return Link{
ID: id, Org: org, User: user, Machine: machine, Host: machine + "-host", OS: "linux",
Provider: provider, Account: account, Kind: kind, Status: StatusLinked,
LastSeen: now, CreatedAt: now, UpdatedAt: now,
}
}
// TestUpsertIdempotent: a repeat report of the same identity updates in place,
// keeps the original id + created_at, refreshes labels, merges usage (an empty
// heartbeat keeps the last snapshot), and re-links a revoked account.
func TestUpsertIdempotent(t *testing.T) {
s := testStore(t)
ctx := context.Background()
a := mkLink("acme", "alice", "m1", "claude", "alice@x", KindSubscription)
a.Usage = `{"sessionPct":10}`
first, err := s.Upsert(ctx, a)
if err != nil {
t.Fatalf("upsert 1: %v", err)
}
// A second report of the SAME identity with a different id/host/usage.
b := mkLink("acme", "alice", "m1", "claude", "alice@x", KindSubscription)
b.Host = "renamed"
b.Usage = `{"sessionPct":80}`
second, err := s.Upsert(ctx, b)
if err != nil {
t.Fatalf("upsert 2: %v", err)
}
if second.ID != first.ID {
t.Fatalf("upsert must keep the original id: %s vs %s", second.ID, first.ID)
}
if second.CreatedAt != first.CreatedAt {
t.Fatalf("upsert must preserve created_at")
}
if second.Host != "renamed" {
t.Fatalf("upsert must refresh the device label, got %q", second.Host)
}
if second.Usage != `{"sessionPct":80}` {
t.Fatalf("upsert must replace usage with the fresh snapshot, got %q", second.Usage)
}
if got, _ := s.List(ctx, "acme", "alice"); len(got) != 1 {
t.Fatalf("upsert of the same identity must NOT dup: want 1 row, got %d", len(got))
}
// A heartbeat with no usage keeps the last good snapshot (keep-stale rule).
c := mkLink("acme", "alice", "m1", "claude", "alice@x", KindSubscription)
c.Usage = ""
third, err := s.Upsert(ctx, c)
if err != nil {
t.Fatalf("upsert 3: %v", err)
}
if third.Usage != `{"sessionPct":80}` {
t.Fatalf("empty-usage heartbeat must keep the last snapshot, got %q", third.Usage)
}
// Revoke, then re-report → re-linked (log back in).
if _, found, err := s.Revoke(ctx, "acme", "alice", first.ID, time.Now().Unix()); err != nil || !found {
t.Fatalf("revoke: found=%v err=%v", found, err)
}
relink, err := s.Upsert(ctx, mkLink("acme", "alice", "m1", "claude", "alice@x", KindSubscription))
if err != nil {
t.Fatalf("re-report: %v", err)
}
if relink.Status != StatusLinked {
t.Fatalf("a re-report must re-link a revoked account, got %q", relink.Status)
}
}
// TestOrgAndUserIsolation: (org, subject) scopes EVERY read/write. A foreign org,
// AND a foreign user in the SAME org, sees/gets/revokes nothing of another's.
func TestOrgAndUserIsolation(t *testing.T) {
s := testStore(t)
ctx := context.Background()
al, _ := s.Upsert(ctx, mkLink("acme", "alice", "m1", "claude", "a", KindSubscription))
_, _ = s.Upsert(ctx, mkLink("acme", "bob", "m2", "codex", "b", KindSubscription)) // same org, other user
_, _ = s.Upsert(ctx, mkLink("evil", "alice", "m1", "claude", "a", KindSubscription))
aliceLinks, _ := s.List(ctx, "acme", "alice")
if len(aliceLinks) != 1 || aliceLinks[0].ID != al.ID {
t.Fatalf("alice must see only her own link, got %+v", aliceLinks)
}
// Alice cannot Get bob's link by its id (same org, different user).
bobLinks, _ := s.List(ctx, "acme", "bob")
if len(bobLinks) != 1 {
t.Fatalf("bob setup: want 1, got %d", len(bobLinks))
}
if _, err := s.Get(ctx, "acme", "alice", bobLinks[0].ID); err != errNotFound {
t.Fatalf("alice must not Get bob's link: err=%v", err)
}
// The evil-org row of the same (user, machine, provider, account) is DISTINCT
// (org is part of the identity) — no cross-org collision.
evilLinks, _ := s.List(ctx, "evil", "alice")
if len(evilLinks) != 1 || evilLinks[0].ID == al.ID {
t.Fatalf("evil-org link must be a distinct row, got %+v", evilLinks)
}
if _, err := s.Get(ctx, "evil", "alice", al.ID); err != errNotFound {
t.Fatalf("alice's acme link must not resolve under evil org")
}
// Bob revoking alice's id is a no-op (found=false) — never a cross-user write.
if _, found, _ := s.Revoke(ctx, "acme", "bob", al.ID, time.Now().Unix()); found {
t.Fatalf("bob must not be able to revoke alice's link")
}
if got, _ := s.Get(ctx, "acme", "alice", al.ID); got.Status != StatusLinked {
t.Fatalf("alice's link must remain linked after bob's attempt, got %q", got.Status)
}
}
// TestRevokeDeviceScoped: revoking a device revokes exactly that machine's linked
// accounts (within the caller's scope), leaves other machines untouched, and is
// idempotent.
func TestRevokeDeviceScoped(t *testing.T) {
s := testStore(t)
ctx := context.Background()
_, _ = s.Upsert(ctx, mkLink("acme", "alice", "m1", "claude", "a1", KindSubscription))
_, _ = s.Upsert(ctx, mkLink("acme", "alice", "m1", "codex", "a2", KindSubscription))
_, _ = s.Upsert(ctx, mkLink("acme", "alice", "m2", "claude", "a3", KindSubscription))
if d, _ := s.ListDevice(ctx, "acme", "alice", "m1"); len(d) != 2 {
t.Fatalf("device m1 want 2 accounts, got %d", len(d))
}
rev, err := s.RevokeDevice(ctx, "acme", "alice", "m1", time.Now().Unix())
if err != nil || len(rev) != 2 {
t.Fatalf("revoke device m1 want 2 revoked, got %d (%v)", len(rev), err)
}
linked, _ := s.ListLinked(ctx, "acme", "alice")
if len(linked) != 1 || linked[0].Machine != "m2" {
t.Fatalf("only m2 stays linked, got %+v", linked)
}
// Idempotent: nothing left linked on m1.
if rev2, _ := s.RevokeDevice(ctx, "acme", "alice", "m1", time.Now().Unix()); len(rev2) != 0 {
t.Fatalf("re-revoke device want 0, got %d", len(rev2))
}
// A foreign user cannot revoke alice's device.
if rev3, _ := s.RevokeDevice(ctx, "acme", "bob", "m2", time.Now().Unix()); len(rev3) != 0 {
t.Fatalf("bob revoking alice's device must be a no-op, got %d", len(rev3))
}
}
+4
View File
@@ -87,6 +87,7 @@ import (
"github.com/hanzoai/cloud/clients/kafka"
"github.com/hanzoai/cloud/clients/kms"
"github.com/hanzoai/cloud/clients/knowledge"
"github.com/hanzoai/cloud/clients/link"
"github.com/hanzoai/cloud/clients/marketing"
"github.com/hanzoai/cloud/clients/marketplace"
"github.com/hanzoai/cloud/clients/ml"
@@ -226,6 +227,9 @@ func Wire() []cloud.MountSpec {
{Name: "projects", Mount: cloud.Typed(projects.Mount)},
{Name: "prompts", Mount: cloud.Typed(prompts.Mount)},
{Name: "agents", Mount: cloud.Typed(agents.Mount), Shutdown: agents.Shutdown},
// The unified AI login manager registry (/v1/links). Mounts AFTER agents so
// a link revoke can stop the affected agent sessions in-process.
{Name: "link", Mount: cloud.Typed(link.Mount), Shutdown: link.Shutdown},
{Name: "wallets", Mount: cloud.Typed(wallets.Mount), Shutdown: ctxShutdown(wallets.Shutdown)},
// x402 pay-per-use: settles a signed ERC-3009 authorization to a recipient
// wallet through the metering spine. Mounts AFTER wallets (it resolves the