Merge pull request #1110 from Molecule-AI/staging

promote: org-tokens review followups

workspace-server/internal/handlers/org_tokens.go

@@ -112,27 +112,37 @@ func (h *OrgTokenHandler) Revoke(c *gin.Context) {
 	c.JSON(http.StatusOK, gin.H{"revoked": id})
 }
 
+// Provenance labels used in the org_api_tokens.created_by column
+// and in mint/revoke audit logs. Kept as constants so the labels
+// are greppable across services (log pipelines, audit queries).
+const (
+	actorOrgTokenPrefix = "org-token:" // appended: 8-char plaintext prefix from the UI
+	actorSession        = "session"    // WorkOS-session-verified call
+	actorAdminToken     = "admin-token" // bootstrap ADMIN_TOKEN env
+)
+
 // orgTokenActor derives a short provenance string for audit.
 //
 //   - If the request was authed via another org token, return
-//     "org-token:<short>" so revoke audits show which token minted
-//     the new one.
+//     "org-token:<prefix>" where prefix is the 8-char plaintext
+//     prefix shown in the UI — correlates audit rows directly with
+//     the revoke button a user sees.
 //   - If authed via session cookie (AdminAuth's session tier), the
-//     middleware doesn't set anything on c for us — return "session"
-//     as a generic label. When we grow a session-user-id capture
-//     upgrade this to return the real WorkOS user_id.
+//     middleware doesn't stash a WorkOS user_id today — return
+//     "session" as a generic label. Follow-up (see
+//     docs/architecture/org-api-keys-followups.md #6) captures the
+//     user_id through the session tier for full attribution.
 //   - Else (ADMIN_TOKEN / bootstrap), return "admin-token".
 func orgTokenActor(c *gin.Context) string {
-	if v, ok := c.Get("org_token_id"); ok {
-		if s, ok := v.(string); ok && len(s) >= 8 {
-			return "org-token:" + s[:8]
+	if v, ok := c.Get("org_token_prefix"); ok {
+		if s, ok := v.(string); ok && s != "" {
+			return actorOrgTokenPrefix + s
 		}
 	}
 	// Session-tier auth doesn't stash an identity in the gin context
-	// today. Until it does, treat session requests as "session". A
-	// follow-up issue captures WorkOS user_id here.
+	// today. Until it does, treat session requests as "session".
 	if c.GetHeader("Cookie") != "" {
-		return "session"
+		return actorSession
 	}
-	return "admin-token"
+	return actorAdminToken
 }

workspace-server/internal/handlers/org_tokens_test.go

@@ -0,0 +1,287 @@
+package handlers
+
+import (
+	"bytes"
+	"context"
+	"database/sql"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
+	"github.com/gin-gonic/gin"
+)
+
+// setupOrgTokenTest wires the package-global db.DB to a sqlmock for
+// the duration of a test, returning the handler + mock + cleanup.
+// Gin runs in release mode to suppress debug noise.
+func setupOrgTokenTest(t *testing.T) (*OrgTokenHandler, sqlmock.Sqlmock, func()) {
+	t.Helper()
+	gin.SetMode(gin.ReleaseMode)
+	mock, mockDB, cleanup := mockSQLDB(t)
+	prev := db.DB
+	db.DB = mockDB
+	return NewOrgTokenHandler(), mock, func() {
+		db.DB = prev
+		cleanup()
+	}
+}
+
+// mockSQLDB returns an sqlmock + *sql.DB pair. Caller restores
+// package state via the cleanup func.
+func mockSQLDB(t *testing.T) (sqlmock.Sqlmock, *sql.DB, func()) {
+	t.Helper()
+	d, m, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("sqlmock: %v", err)
+	}
+	return m, d, func() { _ = d.Close() }
+}
+
+// buildCtx returns a gin.Context + recorder wired for the given
+// method+path+body. Test code can set headers / context values
+// before calling the handler.
+func buildCtx(method, path, body string) (*gin.Context, *httptest.ResponseRecorder) {
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	var r *http.Request
+	if body != "" {
+		r = httptest.NewRequest(method, path, bytes.NewBufferString(body))
+		r.Header.Set("Content-Type", "application/json")
+	} else {
+		r = httptest.NewRequest(method, path, nil)
+	}
+	c.Request = r.WithContext(context.Background())
+	return c, w
+}
+
+// ---- List -----------------------------------------------------------------
+
+func TestOrgTokenHandler_List_HappyPath(t *testing.T) {
+	h, mock, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+
+	now := time.Now().UTC()
+	mock.ExpectQuery(`SELECT id, prefix.*FROM org_api_tokens`).
+		WillReturnRows(sqlmock.NewRows([]string{"id", "prefix", "name", "created_by", "created_at", "last_used_at"}).
+			AddRow("tok-1", "abcd1234", "zapier", "session", now, nil))
+
+	c, w := buildCtx("GET", "/org/tokens", "")
+	h.List(c)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("want 200, got %d: %s", w.Code, w.Body.String())
+	}
+	var body struct {
+		Count  int `json:"count"`
+		Tokens []struct {
+			ID     string `json:"id"`
+			Prefix string `json:"prefix"`
+			Name   string `json:"name"`
+		} `json:"tokens"`
+	}
+	if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
+		t.Fatalf("parse: %v", err)
+	}
+	if body.Count != 1 || len(body.Tokens) != 1 {
+		t.Fatalf("expected 1 token, got %+v", body)
+	}
+	if body.Tokens[0].Prefix != "abcd1234" {
+		t.Errorf("prefix not propagated: %q", body.Tokens[0].Prefix)
+	}
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet: %v", err)
+	}
+}
+
+func TestOrgTokenHandler_List_DBError500(t *testing.T) {
+	h, mock, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+	mock.ExpectQuery(`SELECT id, prefix.*FROM org_api_tokens`).
+		WillReturnError(sql.ErrConnDone)
+
+	c, w := buildCtx("GET", "/org/tokens", "")
+	h.List(c)
+
+	if w.Code != http.StatusInternalServerError {
+		t.Errorf("db error → 500 expected; got %d", w.Code)
+	}
+}
+
+// ---- Create ---------------------------------------------------------------
+
+func TestOrgTokenHandler_Create_ActorFromAdminToken(t *testing.T) {
+	h, mock, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+
+	// No Cookie header, no org_token_prefix → actor should be
+	// "admin-token" (bootstrap path).
+	mock.ExpectQuery(`INSERT INTO org_api_tokens`).
+		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), "my-ci", actorAdminToken).
+		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-1"))
+
+	c, w := buildCtx("POST", "/org/tokens", `{"name":"my-ci"}`)
+	h.Create(c)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("want 200, got %d: %s", w.Code, w.Body.String())
+	}
+	var body struct {
+		ID      string `json:"id"`
+		Prefix  string `json:"prefix"`
+		Name    string `json:"name"`
+		Token   string `json:"auth_token"`
+		Warning string `json:"warning"`
+	}
+	if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
+		t.Fatalf("parse: %v", err)
+	}
+	if body.Token == "" {
+		t.Errorf("plaintext auth_token missing from response")
+	}
+	if body.Prefix != body.Token[:8] {
+		t.Errorf("prefix %q should equal first 8 chars of token %q", body.Prefix, body.Token[:8])
+	}
+	if body.Name != "my-ci" {
+		t.Errorf("name round-trip mismatch: %q", body.Name)
+	}
+	if body.Warning == "" {
+		t.Errorf("warning text missing")
+	}
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet: %v", err)
+	}
+}
+
+func TestOrgTokenHandler_Create_ActorFromOrgTokenPrefix(t *testing.T) {
+	h, mock, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+
+	// When an existing org token authenticated the mint, audit
+	// records "org-token:<prefix>" using the 8-char plaintext
+	// prefix from the presenter's token.
+	mock.ExpectQuery(`INSERT INTO org_api_tokens`).
+		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), nil, actorOrgTokenPrefix+"parent12").
+		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-new"))
+
+	c, w := buildCtx("POST", "/org/tokens", `{}`)
+	c.Set("org_token_prefix", "parent12")
+	h.Create(c)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("want 200, got %d: %s", w.Code, w.Body.String())
+	}
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet: %v", err)
+	}
+}
+
+func TestOrgTokenHandler_Create_ActorFromSession(t *testing.T) {
+	h, mock, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+
+	// Cookie present but no org_token_prefix — that's the canvas
+	// session path. Today recorded as bare "session". When the
+	// follow-up captures WorkOS user_id, this test changes to
+	// "session:user_01XXX".
+	mock.ExpectQuery(`INSERT INTO org_api_tokens`).
+		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), "from-browser", actorSession).
+		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-browser"))
+
+	c, w := buildCtx("POST", "/org/tokens", `{"name":"from-browser"}`)
+	c.Request.Header.Set("Cookie", "mcp_session=abc")
+	h.Create(c)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("want 200, got %d: %s", w.Code, w.Body.String())
+	}
+}
+
+func TestOrgTokenHandler_Create_NameTooLong400(t *testing.T) {
+	h, _, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+	longName := string(make([]byte, 101))
+	for i := range longName {
+		_ = i
+	}
+	// Build a 101-char name (any ASCII works; we're hitting the
+	// length bound).
+	buf := make([]byte, 101)
+	for i := range buf {
+		buf[i] = 'a'
+	}
+	c, w := buildCtx("POST", "/org/tokens", `{"name":"`+string(buf)+`"}`)
+	h.Create(c)
+
+	if w.Code != http.StatusBadRequest {
+		t.Errorf("oversize name: want 400, got %d", w.Code)
+	}
+}
+
+func TestOrgTokenHandler_Create_EmptyBodyOK(t *testing.T) {
+	h, mock, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+	// Empty POST must still mint a token — name is optional.
+	mock.ExpectQuery(`INSERT INTO org_api_tokens`).
+		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), nil, actorAdminToken).
+		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-min"))
+
+	c, w := buildCtx("POST", "/org/tokens", "")
+	h.Create(c)
+
+	if w.Code != http.StatusOK {
+		t.Errorf("empty body should mint OK; got %d: %s", w.Code, w.Body.String())
+	}
+}
+
+// ---- Revoke ---------------------------------------------------------------
+
+func TestOrgTokenHandler_Revoke_HappyPath200(t *testing.T) {
+	h, mock, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+
+	mock.ExpectExec(`UPDATE org_api_tokens`).
+		WithArgs("tok-1").
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	c, w := buildCtx("DELETE", "/org/tokens/tok-1", "")
+	c.Params = gin.Params{{Key: "id", Value: "tok-1"}}
+	h.Revoke(c)
+
+	if w.Code != http.StatusOK {
+		t.Errorf("want 200, got %d: %s", w.Code, w.Body.String())
+	}
+}
+
+func TestOrgTokenHandler_Revoke_Missing404(t *testing.T) {
+	// Idempotency: revoking a non-existent or already-revoked id
+	// returns 404 — callers can tell "worked" from "already done".
+	h, mock, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+	mock.ExpectExec(`UPDATE org_api_tokens`).
+		WithArgs("ghost").
+		WillReturnResult(sqlmock.NewResult(0, 0))
+
+	c, w := buildCtx("DELETE", "/org/tokens/ghost", "")
+	c.Params = gin.Params{{Key: "id", Value: "ghost"}}
+	h.Revoke(c)
+
+	if w.Code != http.StatusNotFound {
+		t.Errorf("want 404, got %d", w.Code)
+	}
+}
+
+func TestOrgTokenHandler_Revoke_MissingID400(t *testing.T) {
+	h, _, cleanup := setupOrgTokenTest(t)
+	defer cleanup()
+	c, w := buildCtx("DELETE", "/org/tokens/", "")
+	// No id param set.
+	h.Revoke(c)
+
+	if w.Code != http.StatusBadRequest {
+		t.Errorf("want 400, got %d", w.Code)
+	}
+}

workspace-server/internal/middleware/wsauth_middleware.go

@@ -60,8 +60,9 @@ func WorkspaceAuth(database *sql.DB) gin.HandlerFunc {
 			// power surface as ADMIN_TOKEN but named, revocable, audited.
 			// Check before per-workspace token so an org-key presenter
 			// doesn't hit the narrower ValidateToken failure path.
-			if id, err := orgtoken.Validate(ctx, database, tok); err == nil {
+			if id, prefix, err := orgtoken.Validate(ctx, database, tok); err == nil {
 				c.Set("org_token_id", id)
+				c.Set("org_token_prefix", prefix)
 				c.Next()
 				return
 			} else if !errors.Is(err, orgtoken.ErrInvalidToken) {
@@ -180,8 +181,9 @@ func AdminAuth(database *sql.DB) gin.HandlerFunc {
 		// index with revoked_at IS NULL) + an async last_used_at
 		// bump. Cost per request: one SELECT + one UPDATE, both
 		// hitting the same narrow partial index.
-		if id, err := orgtoken.Validate(ctx, database, tok); err == nil {
+		if id, prefix, err := orgtoken.Validate(ctx, database, tok); err == nil {
 			c.Set("org_token_id", id)
+			c.Set("org_token_prefix", prefix)
 			c.Next()
 			return
 		} else if !errors.Is(err, orgtoken.ErrInvalidToken) {

workspace-server/internal/orgtoken/tokens.go

@@ -36,6 +36,12 @@ const (
 	// crackable on its own (6 bits × 8 = 48 bits of prefix space —
 	// good enough to disambiguate, nowhere near guessable).
 	tokenPrefixLen = 8
+
+	// listMax caps the number of rows List returns. Realistic org-
+	// admin UIs show on the order of 10. 500 is enough headroom that
+	// no legitimate flow hits it, low enough that a mint-storm can't
+	// force a big allocation on every list render.
+	listMax = 500
 )
 
 // ErrInvalidToken is returned when a presented bearer doesn't match
@@ -86,33 +92,42 @@ func Issue(ctx context.Context, db *sql.DB, name, createdBy string) (plaintext,
 // Validate looks up a presented bearer, returns ErrInvalidToken on
 // any mismatch (bad bytes, revoked, deleted). On success, updates
 // last_used_at best-effort (the hot path — failure to update doesn't
-// fail the request) and returns the token id for audit logging.
-func Validate(ctx context.Context, db *sql.DB, plaintext string) (string, error) {
+// fail the request) and returns the token id + display prefix for
+// audit logging.
+//
+// Returning the prefix alongside the id lets callers produce audit
+// strings that match what users see in the UI (the plaintext prefix,
+// not the UUID). Keeps the "who did what" trail visually
+// correlatable to the revoke button in the token list.
+func Validate(ctx context.Context, db *sql.DB, plaintext string) (id, prefix string, err error) {
 	if plaintext == "" {
-		return "", ErrInvalidToken
+		return "", "", ErrInvalidToken
 	}
 	hash := sha256.Sum256([]byte(plaintext))
-	var id string
-	err := db.QueryRowContext(ctx, `
-		SELECT id FROM org_api_tokens
+	queryErr := db.QueryRowContext(ctx, `
+		SELECT id, prefix FROM org_api_tokens
 		WHERE token_hash = $1 AND revoked_at IS NULL
-	`, hash[:]).Scan(&id)
-	if err != nil {
+	`, hash[:]).Scan(&id, &prefix)
+	if queryErr != nil {
 		// Collapse all failure shapes into ErrInvalidToken so the
 		// caller can't accidentally leak "row exists but revoked" vs
 		// "row never existed" via response shape.
-		return "", ErrInvalidToken
+		return "", "", ErrInvalidToken
 	}
 	// Best-effort last_used_at bump. Failure here is acceptable — the
 	// request is already authenticated; we don't want a transient DB
 	// blip to flip a 200 into a 500.
 	_, _ = db.ExecContext(ctx,
 		`UPDATE org_api_tokens SET last_used_at = now() WHERE id = $1`, id)
-	return id, nil
+	return id, prefix, nil
 }
 
 // List returns live (non-revoked) tokens newest-first. Safe to
 // expose to the admin UI — no hash, no plaintext, only prefix.
+//
+// Capped at listMax rows. A UI page with more than that is a
+// symptom of abuse or a bug — the hard cap prevents one runaway
+// minting loop from O(N) pageloads in the admin UI.
 func List(ctx context.Context, db *sql.DB) ([]Token, error) {
 	rows, err := db.QueryContext(ctx, `
 		SELECT id, prefix, COALESCE(name,''), COALESCE(created_by,''),
@@ -120,7 +135,8 @@ func List(ctx context.Context, db *sql.DB) ([]Token, error) {
 		FROM org_api_tokens
 		WHERE revoked_at IS NULL
 		ORDER BY created_at DESC
-	`)
+		LIMIT $1
+	`, listMax)
 	if err != nil {
 		return nil, fmt.Errorf("orgtoken: list: %w", err)
 	}

workspace-server/internal/orgtoken/tokens_test.go

@@ -72,26 +72,29 @@ func TestValidate_HappyPath(t *testing.T) {
 	plaintext := "known-plaintext-for-test"
 	hash := sha256.Sum256([]byte(plaintext))
 
-	mock.ExpectQuery(`SELECT id FROM org_api_tokens`).
+	mock.ExpectQuery(`SELECT id, prefix FROM org_api_tokens`).
 		WithArgs(hash[:]).
-		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-live"))
+		WillReturnRows(sqlmock.NewRows([]string{"id", "prefix"}).AddRow("tok-live", "abcd1234"))
 	mock.ExpectExec(`UPDATE org_api_tokens SET last_used_at`).
 		WithArgs("tok-live").
 		WillReturnResult(sqlmock.NewResult(0, 1))
 
-	id, err := Validate(context.Background(), db, plaintext)
+	id, prefix, err := Validate(context.Background(), db, plaintext)
 	if err != nil {
 		t.Fatalf("Validate: %v", err)
 	}
 	if id != "tok-live" {
 		t.Errorf("id = %q, want tok-live", id)
 	}
+	if prefix != "abcd1234" {
+		t.Errorf("prefix = %q, want abcd1234", prefix)
+	}
 }
 
 func TestValidate_EmptyPlaintextRejected(t *testing.T) {
 	db, _, _ := sqlmock.New()
 	defer db.Close()
-	if _, err := Validate(context.Background(), db, ""); !errors.Is(err, ErrInvalidToken) {
+	if _, _, err := Validate(context.Background(), db, ""); !errors.Is(err, ErrInvalidToken) {
 		t.Errorf("empty plaintext should be ErrInvalidToken, got %v", err)
 	}
 }
@@ -103,11 +106,11 @@ func TestValidate_UnknownHashErrInvalid(t *testing.T) {
 	}
 	defer db.Close()
 
-	mock.ExpectQuery(`SELECT id FROM org_api_tokens`).
+	mock.ExpectQuery(`SELECT id, prefix FROM org_api_tokens`).
 		WithArgs(sqlmock.AnyArg()).
 		WillReturnError(sql.ErrNoRows)
 
-	if _, err := Validate(context.Background(), db, "ghost"); !errors.Is(err, ErrInvalidToken) {
+	if _, _, err := Validate(context.Background(), db, "ghost"); !errors.Is(err, ErrInvalidToken) {
 		t.Errorf("unknown hash should be ErrInvalidToken, got %v", err)
 	}
 }
@@ -120,11 +123,11 @@ func TestValidate_RevokedTokenNotAccepted(t *testing.T) {
 	defer db.Close()
 	// Query has `AND revoked_at IS NULL` — sqlmock will return
 	// ErrNoRows because the revoked row is filtered out.
-	mock.ExpectQuery(`SELECT id FROM org_api_tokens`).
+	mock.ExpectQuery(`SELECT id, prefix FROM org_api_tokens`).
 		WithArgs(sqlmock.AnyArg()).
 		WillReturnError(sql.ErrNoRows)
 
-	if _, err := Validate(context.Background(), db, "revoked-plaintext"); !errors.Is(err, ErrInvalidToken) {
+	if _, _, err := Validate(context.Background(), db, "revoked-plaintext"); !errors.Is(err, ErrInvalidToken) {
 		t.Errorf("revoked token should be ErrInvalidToken, got %v", err)
 	}
 }
@@ -139,6 +142,7 @@ func TestList_NewestFirst(t *testing.T) {
 	now := time.Now()
 	earlier := now.Add(-1 * time.Hour)
 	mock.ExpectQuery(`SELECT id, prefix.*FROM org_api_tokens.*ORDER BY created_at DESC`).
+		WithArgs(listMax).
 		WillReturnRows(sqlmock.NewRows([]string{"id", "prefix", "name", "created_by", "created_at", "last_used_at"}).
 			AddRow("t2", "abcd1234", "zapier", "user_01", now, now).
 			AddRow("t1", "efgh5678", "", "", earlier, nil))

workspace-server/internal/router/router.go

@@ -476,11 +476,20 @@ func Setup(hub *ws.Hub, broadcaster *events.Broadcaster, prov *provisioner.Provi
 	// session cookie, OR an existing org token to mint more. That's
 	// bootstrap-friendly (first token from ADMIN_TOKEN or canvas
 	// session) and self-sustaining afterwards (tokens mint tokens).
+	//
+	// The mint endpoint gets an extra per-IP rate limiter — a
+	// compromised session or leaked bearer could otherwise mint
+	// thousands of tokens per second, making forensic cleanup
+	// painful. 10 mints per hour per IP is ample for real usage;
+	// legitimate bursts fit in the ceiling and abuse bounces off.
+	// List + Delete don't need the extra limit (they can't be used
+	// to generate new secret material).
 	{
 		orgTokenHandler := handlers.NewOrgTokenHandler()
 		orgTokenAdmin := r.Group("", middleware.AdminAuth(db.DB))
 		orgTokenAdmin.GET("/org/tokens", orgTokenHandler.List)
-		orgTokenAdmin.POST("/org/tokens", orgTokenHandler.Create)
+		orgTokenMintLimiter := middleware.NewRateLimiter(10, time.Hour, context.Background())
+		orgTokenAdmin.POST("/org/tokens", orgTokenMintLimiter.Middleware(), orgTokenHandler.Create)
 		orgTokenAdmin.DELETE("/org/tokens/:id", orgTokenHandler.Revoke)
 	}