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:
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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) }
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user