core-lead/molecule-core
1
0
Fork 0
forked from molecule-ai/molecule-core
Code Pull Requests Activity

fix(team): delegate Expand child-provisioning to shared mint pipeline (#2367)

Browse Source 
Closes #2367.

TeamHandler.Expand provisioned child workspaces by directly calling
h.provisioner.Start, skipping mintWorkspaceSecrets and every other
preflight (secrets load, env mutators, identity injection, missing-env,
empty-config-volume auto-recover). Children shipped with NULL
platform_inbound_secret + never-issued auth_token — same drift class as
the SaaS bug just fixed in PR #2366, found while exercising a stronger
gate against this package.

Fix:

- TeamHandler now holds *WorkspaceHandler. Expand delegates each child
  provision to wh.provisionWorkspace, picking up the shared
  prepare/mint/preflight pipeline automatically. Future provision-time
  steps go in ONE place and team-expand inherits them.
- prepareProvisionContext gains PARENT_ID env injection sourced from
  payload.ParentID (which Expand now populates). This preserves the
  signal workspace/coordinator.py reads on startup, without threading
  env through provisioner.WorkspaceConfig manually.
- NewTeamHandler signature gains *WorkspaceHandler; router passes it.

Gate upgrade:

- TestProvisionFunctions_AllCallMintWorkspaceSecrets is now
  behavior-based: it walks every FuncDecl in the package and flags any
  function that calls h.provisioner.Start or h.cpProv.Start without
  also calling mintWorkspaceSecrets. Drift-resistant by construction —
  a future provision function with any name still trips the gate.
- Replaces the name-list version from PR #2366. The name list missed
  Expand precisely because Expand wasn't named provision*; the
  behavior-based detector caught it spontaneously when prototyped.

Tests: full workspace-server module green; gate previously verified to
fire red on Expand pre-fix and on deliberate mintWorkspaceSecrets
removal.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Hongming Wang 2026-04-30 02:28:12 -07:00
parent b9b0a46f2e
commit bb52a1a365
5 changed files with 103 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
workspace-server/internal/handlers/team.go
View File

@ -1,7 +1,6 @@
package handlers
import (
"context"
"database/sql"
"encoding/json"
"log"
@ -11,23 +10,32 @@ import (
"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
"github.com/Molecule-AI/molecule-monorepo/platform/internal/events"
"github.com/Molecule-AI/molecule-monorepo/platform/internal/models"
"github.com/Molecule-AI/molecule-monorepo/platform/internal/provisioner"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"gopkg.in/yaml.v3"
)
// TeamHandler delegates child-workspace provisioning to wh so child
// workspaces go through the same prepare/mint/preflight pipeline that
// every other provision path uses. Pre-fix (issue #2367): Expand
// directly invoked h.provisioner.Start which skipped mintWorkspaceSecrets,
// leaving every team-expanded child with NULL platform_inbound_secret +
// NULL auth_token — same drift class as the SaaS bug fixed in #2366.
type TeamHandler struct {
broadcaster *events.Broadcaster
provisioner *provisioner.Provisioner
wh *WorkspaceHandler
platformURL string
configsDir string
}
func NewTeamHandler(b *events.Broadcaster, p *provisioner.Provisioner, platformURL, configsDir string) *TeamHandler {
func NewTeamHandler(b *events.Broadcaster, p *provisioner.Provisioner, wh *WorkspaceHandler, platformURL, configsDir string) *TeamHandler {
return &TeamHandler{
broadcaster: b,
provisioner: p,
wh: wh,
platformURL: platformURL,
configsDir: configsDir,
}
@ -116,26 +124,25 @@ func (h *TeamHandler) Expand(c *gin.Context) {
"parent_id": parentID,
})
// Provision if template exists
if h.provisioner != nil && sub.Config != "" {
// Delegate child-workspace provisioning to the shared
// provision pipeline. Issue #2367: previously Expand called
// h.provisioner.Start directly, bypassing mintWorkspaceSecrets
// and every other preflight (secrets, env mutators, identity
// injection, missing-env). That left every child with NULL
// platform_inbound_secret and never-issued auth_token. Now
// children go through the same provisionWorkspace path as
// Create/Restart, so adding a future provision-time step
// automatically covers Expand too.
if h.wh != nil && sub.Config != "" {
templatePath := filepath.Join(h.configsDir, sub.Config)
if _, err := os.Stat(templatePath); err == nil {
pluginsPath, _ := filepath.Abs(filepath.Join(h.configsDir, "..", "plugins"))
go func(wID, tPath, pPath string, t int) {
provCtx, cancel := context.WithTimeout(context.Background(), provisioner.ProvisionTimeout)
defer cancel()
cfg := provisioner.WorkspaceConfig{
WorkspaceID: wID,
TemplatePath: tPath,
PluginsPath: pPath,
Tier: t,
EnvVars: map[string]string{"PARENT_ID": parentID},
PlatformURL: h.platformURL,
}
if _, err := h.provisioner.Start(provCtx, cfg); err != nil {
log.Printf("Expand: provision failed for %s: %v", wID, err)
}
}(childID, templatePath, pluginsPath, tier)
parent := parentID // copy for closure
go h.wh.provisionWorkspace(childID, templatePath, nil, models.CreateWorkspacePayload{
Name: childName,
Role: sub.Role,
Tier: tier,
ParentID: &parent,
})
}
}

12
workspace-server/internal/handlers/team_test.go
View File

@ -34,7 +34,7 @@ func TestTeamCollapse_NoChildren(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
broadcaster := newTestBroadcaster()
handler := NewTeamHandler(broadcaster, nil, "http://localhost:8080", "/tmp/configs")
handler := NewTeamHandler(broadcaster, nil, nil, "http://localhost:8080", "/tmp/configs")
// No children
mock.ExpectQuery("SELECT id, name FROM workspaces WHERE parent_id").
@ -66,7 +66,7 @@ func TestTeamCollapse_WithChildren(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
broadcaster := newTestBroadcaster()
handler := NewTeamHandler(broadcaster, nil, "http://localhost:8080", "/tmp/configs")
handler := NewTeamHandler(broadcaster, nil, nil, "http://localhost:8080", "/tmp/configs")
// Two children
mock.ExpectQuery("SELECT id, name FROM workspaces WHERE parent_id").
@ -122,7 +122,7 @@ func TestTeamCollapse_WithChildren(t *testing.T) {
func TestTeamExpand_WorkspaceNotFound(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
handler := NewTeamHandler(newTestBroadcaster(), nil, "http://localhost:8080", "/tmp/configs")
handler := NewTeamHandler(newTestBroadcaster(), nil, nil, "http://localhost:8080", "/tmp/configs")
mock.ExpectQuery("SELECT name, tier, status FROM workspaces WHERE id").
WithArgs("ws-missing").
@ -143,7 +143,7 @@ func TestTeamExpand_WorkspaceNotFound(t *testing.T) {
func TestTeamExpand_NoConfigFound(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
handler := NewTeamHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
handler := NewTeamHandler(newTestBroadcaster(), nil, nil, "http://localhost:8080", t.TempDir())
mock.ExpectQuery("SELECT name, tier, status FROM workspaces WHERE id").
WithArgs("ws-1").
@ -167,7 +167,7 @@ func TestTeamExpand_EmptySubWorkspaces(t *testing.T) {
setupTestRedis(t)
configDir := makeTeamConfigDir(t, "myagent", "name: MyAgent\nsub_workspaces: []\n")
handler := NewTeamHandler(newTestBroadcaster(), nil, "http://localhost:8080", configDir)
handler := NewTeamHandler(newTestBroadcaster(), nil, nil, "http://localhost:8080", configDir)
mock.ExpectQuery("SELECT name, tier, status FROM workspaces WHERE id").
WithArgs("ws-1").
@ -199,7 +199,7 @@ sub_workspaces:
role: code-reviewer
`
configDir := makeTeamConfigDir(t, "teamlead", yaml)
handler := NewTeamHandler(broadcaster, nil, "http://localhost:8080", configDir)
handler := NewTeamHandler(broadcaster, nil, nil, "http://localhost:8080", configDir)
mock.ExpectQuery("SELECT name, tier, status FROM workspaces WHERE id").
WithArgs("ws-lead").

8
workspace-server/internal/handlers/workspace_provision_shared.go
View File

@ -99,6 +99,14 @@ func (h *WorkspaceHandler) prepareProvisionContext(
if payload.Role != "" {
envVars["MOLECULE_AGENT_ROLE"] = payload.Role
}
// PARENT_ID is consumed by workspace/coordinator.py to track the
// parent-child relationship at runtime. Sourced from payload so
// every provision path that knows about a parent (currently:
// TeamHandler.Expand) injects it without having to thread env
// through provisioner.WorkspaceConfig manually.
if payload.ParentID != nil && *payload.ParentID != "" {
envVars["PARENT_ID"] = *payload.ParentID
}
// Plugin extension point: env mutators run AFTER built-in identity
// injection so plugins can override or augment identity vars.

83
workspace-server/internal/handlers/workspace_provision_shared_test.go
View File

@ -31,31 +31,32 @@ import (
"github.com/Molecule-AI/molecule-monorepo/platform/internal/provisioner"
)
// provisionFunctions are the set of functions that provision a
// workspace and therefore MUST call mintWorkspaceSecrets. Detected
// by name match (handler-method pattern); add a new entry whenever
// a new provision path is introduced.
var provisionFunctions = map[string]bool{
"provisionWorkspace": true,
"provisionWorkspaceOpts": true,
"provisionWorkspaceCP": true,
}
// provisionExemptFunctions are functions whose name pattern matches
// the provision-function set but legitimately don't call
// mintWorkspaceSecrets (e.g., the wrapper provisionWorkspace which
// delegates to provisionWorkspaceOpts — the actual mint happens in
// the delegate, not the wrapper).
// provisionExemptFunctions are functions that call a provision-start
// method but legitimately do NOT need to mint (e.g. the wrapper
// `provisionWorkspace` which delegates — the delegate mints; the
// re-spawn loops inside Restart that re-enter provisionWorkspaceOpts).
// Add an entry only with a one-line justification.
var provisionExemptFunctions = map[string]string{
"provisionWorkspace": "thin wrapper that delegates to provisionWorkspaceOpts; the delegate mints",
}
// TestProvisionFunctions_AllCallMintWorkspaceSecrets asserts every
// non-exempt provision function in this package calls
// mintWorkspaceSecrets at least once in its body.
// function in this package that triggers a workspace provision (i.e.
// calls h.provisioner.Start or h.cpProv.Start) ALSO calls
// mintWorkspaceSecrets at least once in the same body.
//
// The check is presence-in-function, not pairing-by-line — same
// rationale as the audit-coverage gate from #335.
// Behavior-based — drift-resistant. A future provision function with
// any name still trips this gate as long as it calls one of the
// provisioner Start methods. This replaces an earlier name-list
// version (PR #2366) that missed TeamHandler.Expand (issue #2367) —
// the bug that motivated the upgrade.
//
// Same shape as the audit-coverage gate from #335 (#2343 PR-5).
//
// If this test fails: either add mintWorkspaceSecrets to the
// offending function (preferred — usually you should delegate to
// provisionWorkspace via h.wh), OR add it to provisionExemptFunctions
// with a one-line justification.
func TestProvisionFunctions_AllCallMintWorkspaceSecrets(t *testing.T) {
t.Parallel()
@ -87,10 +88,10 @@ func TestProvisionFunctions_AllCallMintWorkspaceSecrets(t *testing.T) {
if !ok || fn.Body == nil {
continue
}
if !provisionFunctions[fn.Name.Name] {
if _, exempt := provisionExemptFunctions[fn.Name.Name]; exempt {
continue
}
if _, exempt := provisionExemptFunctions[fn.Name.Name]; exempt {
if !callsProvisionStart(fn.Body) {
continue
}
if !callsMintWorkspaceSecrets(fn.Body) {
@ -105,12 +106,50 @@ func TestProvisionFunctions_AllCallMintWorkspaceSecrets(t *testing.T) {
for _, v := range violations {
t.Errorf(
"%s:%d %s does not call mintWorkspaceSecrets — every provision path MUST mint auth_token + platform_inbound_secret. If audit is genuinely impossible here, add %q to provisionExemptFunctions in workspace_provision_shared_test.go with a justification.",
"%s:%d %s calls a provisioner Start (h.provisioner.Start or h.cpProv.Start) but does not call mintWorkspaceSecrets — every provision path MUST mint auth_token + platform_inbound_secret. Prefer delegating to h.wh.provisionWorkspace; only add %q to provisionExemptFunctions with a one-line justification if mint is genuinely inappropriate.",
v.file, v.line, v.fn, v.fn,
)
}
}
// callsProvisionStart reports whether the function body invokes a
// provisioner-start method. Matches `<x>.provisioner.Start(...)` and
// `<x>.cpProv.Start(...)` — both look like
// `<recv>.<provField>.Start(...)` in the AST. Filtering on the
// provisioner-field name (`provisioner` or `cpProv`) keeps the gate
// from tripping on unrelated `.Start()` calls (e.g. http.Server.Start
// in the same package).
func callsProvisionStart(body *ast.BlockStmt) bool {
found := false
ast.Inspect(body, func(n ast.Node) bool {
if found {
return false
}
call, ok := n.(*ast.CallExpr)
if !ok {
return true
}
sel, ok := call.Fun.(*ast.SelectorExpr)
if !ok {
return true
}
if sel.Sel.Name != "Start" {
return true
}
inner, ok := sel.X.(*ast.SelectorExpr)
if !ok {
return true
}
switch inner.Sel.Name {
case "provisioner", "cpProv":
found = true
return false
}
return true
})
return found
}
// callsMintWorkspaceSecrets walks the function body and reports
// whether mintWorkspaceSecrets is called anywhere — direct call OR
// via a helper. Recursion to helpers is shallow: we only check

2
workspace-server/internal/router/router.go
View File

@ -213,7 +213,7 @@ func Setup(hub *ws.Hub, broadcaster *events.Broadcaster, prov *provisioner.Provi
r.GET("/approvals/pending", middleware.AdminAuth(db.DB), apph.ListAll)
// Team Expansion
teamh := handlers.NewTeamHandler(broadcaster, prov, platformURL, configsDir)
teamh := handlers.NewTeamHandler(broadcaster, prov, wh, platformURL, configsDir)
wsAuth.POST("/expand", teamh.Expand)
wsAuth.POST("/collapse", teamh.Collapse)