diff --git a/workspace-server/internal/handlers/handlers_additional_test.go b/workspace-server/internal/handlers/handlers_additional_test.go index 0e13600d5..98aa736d8 100644 --- a/workspace-server/internal/handlers/handlers_additional_test.go +++ b/workspace-server/internal/handlers/handlers_additional_test.go @@ -230,7 +230,7 @@ func TestWorkspaceList_WithData(t *testing.T) { broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) - // 23 cols — broadcast_enabled + talk_to_user_enabled added after monthly_spend + // 24 cols — compute added after talk_to_user_enabled. // (migration 20260514). Column order must match scanWorkspaceRow exactly. columns := []string{ "id", "name", "role", "tier", "status", "agent_card", "url", @@ -238,13 +238,13 @@ func TestWorkspaceList_WithData(t *testing.T) { "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } rows := sqlmock.NewRows(columns). AddRow("ws-1", "Agent One", "worker", 1, "online", []byte(`{"name":"agent1"}`), "http://localhost:8001", - nil, 3, 1, 0.02, "", 7200, "processing", "langgraph", "", 10.0, 20.0, false, nil, int64(0), false, true). + nil, 3, 1, 0.02, "", 7200, "processing", "langgraph", "", 10.0, 20.0, false, nil, int64(0), false, true, []byte(`{}`)). AddRow("ws-2", "Agent Two", "", 2, "degraded", []byte("null"), "", - nil, 0, 1, 0.6, "timeout", 100, "", "claude-code", "", 50.0, 60.0, true, nil, int64(0), false, true) + nil, 0, 1, 0.6, "timeout", 100, "", "claude-code", "", 50.0, 60.0, true, nil, int64(0), false, true, []byte(`{}`)) mock.ExpectQuery("SELECT w.id, w.name"). WillReturnRows(rows) diff --git a/workspace-server/internal/handlers/handlers_test.go b/workspace-server/internal/handlers/handlers_test.go index 7ce01b239..d14a249de 100644 --- a/workspace-server/internal/handlers/handlers_test.go +++ b/workspace-server/internal/handlers/handlers_test.go @@ -456,7 +456,7 @@ func TestWorkspaceList(t *testing.T) { broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", "/tmp/configs") - // 23 cols: broadcast_enabled + talk_to_user_enabled added after monthly_spend + // 24 cols: compute added after talk_to_user_enabled. // (migration 20260514). Column order must match scanWorkspaceRow exactly. columns := []string{ "id", "name", "role", "tier", "status", "agent_card", "url", @@ -464,13 +464,13 @@ func TestWorkspaceList(t *testing.T) { "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } rows := sqlmock.NewRows(columns). AddRow("ws-1", "Agent One", "worker", 1, "online", []byte("null"), "http://localhost:8001", - nil, 0, 1, 0.0, "", 100, "", "claude-code", "", 10.0, 20.0, false, nil, int64(0), false, true). + nil, 0, 1, 0.0, "", 100, "", "claude-code", "", 10.0, 20.0, false, nil, int64(0), false, true, []byte(`{}`)). AddRow("ws-2", "Agent Two", "manager", 2, "provisioning", []byte("null"), "", - nil, 0, 1, 0.0, "", 0, "", "langgraph", "", 50.0, 60.0, false, nil, int64(0), false, true) + nil, 0, 1, 0.0, "", 0, "", "langgraph", "", 50.0, 60.0, false, nil, int64(0), false, true, []byte(`{}`)) mock.ExpectQuery("SELECT w.id, w.name"). WillReturnRows(rows) @@ -1184,14 +1184,14 @@ func TestWorkspaceGet_CurrentTask(t *testing.T) { "parent_id", "active_tasks", "max_concurrent_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } mock.ExpectQuery("SELECT w.id, w.name"). WithArgs("dddddddd-0004-0000-0000-000000000000"). WillReturnRows(sqlmock.NewRows(columns).AddRow( "dddddddd-0004-0000-0000-000000000000", "Task Worker", "worker", 1, "online", []byte("null"), "http://localhost:9000", nil, 2, 1, 0.0, "", 300, "Analyzing document", "langgraph", "", 10.0, 20.0, false, - nil, int64(0), false, true, + nil, int64(0), false, true, []byte(`{}`), )) w := httptest.NewRecorder() diff --git a/workspace-server/internal/handlers/workspace.go b/workspace-server/internal/handlers/workspace.go index c89622fde..e36af77aa 100644 --- a/workspace-server/internal/handlers/workspace.go +++ b/workspace-server/internal/handlers/workspace.go @@ -348,6 +348,10 @@ func (h *WorkspaceHandler) Create(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid workspace access"}) return } + if err := validateWorkspaceCompute(payload.Compute); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } // Begin a transaction so the workspace row and any initial secrets are // committed atomically. A secret-encrypt or DB error rolls back the @@ -435,6 +439,24 @@ func (h *WorkspaceHandler) Create(c *gin.Context) { payload.Name = persistedName } + if !workspaceComputeIsZero(payload.Compute) { + computeJSON, encErr := workspaceComputeJSON(payload.Compute) + if encErr != nil { + tx.Rollback() //nolint:errcheck + log.Printf("Create workspace %s: failed to encode compute config: %v", id, encErr) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to encode compute config"}) + return + } + if _, dbErr := tx.ExecContext(ctx, + `UPDATE workspaces SET compute = $2::jsonb, updated_at = now() WHERE id = $1`, + id, computeJSON); dbErr != nil { + tx.Rollback() //nolint:errcheck + log.Printf("Create workspace %s: failed to persist compute config: %v", id, dbErr) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to save compute config"}) + return + } + } + // Persist initial secrets from the create payload (inside same transaction). // nil/empty map is a no-op. Any failure rolls back the workspace insert // so we never have a workspace row without its intended secrets. @@ -679,6 +701,7 @@ func scanWorkspaceRow(rows interface { Scan(dest ...interface{}) error }) (map[string]interface{}, error) { var id, name, role, status, url, sampleError, currentTask, runtime, workspaceDir string + var computeRaw []byte var tier, activeTasks, maxConcurrentTasks, uptimeSeconds int var errorRate, x, y float64 var collapsed, broadcastEnabled, talkToUserEnabled bool @@ -690,7 +713,7 @@ func scanWorkspaceRow(rows interface { err := rows.Scan(&id, &name, &role, &tier, &status, &agentCard, &url, &parentID, &activeTasks, &maxConcurrentTasks, &errorRate, &sampleError, &uptimeSeconds, ¤tTask, &runtime, &workspaceDir, &x, &y, &collapsed, - &budgetLimit, &monthlySpend, &broadcastEnabled, &talkToUserEnabled) + &budgetLimit, &monthlySpend, &broadcastEnabled, &talkToUserEnabled, &computeRaw) if err != nil { return nil, err } @@ -717,6 +740,11 @@ func scanWorkspaceRow(rows interface { "broadcast_enabled": broadcastEnabled, "talk_to_user_enabled": talkToUserEnabled, } + if len(computeRaw) > 0 && string(computeRaw) != "null" { + ws["compute"] = json.RawMessage(computeRaw) + } else { + ws["compute"] = json.RawMessage(`{}`) + } // budget_limit: nil when no limit set, int64 otherwise if budgetLimit.Valid { @@ -752,7 +780,8 @@ const workspaceListQuery = ` COALESCE(w.workspace_dir, ''), COALESCE(cl.x, 0), COALESCE(cl.y, 0), COALESCE(cl.collapsed, false), w.budget_limit, COALESCE(w.monthly_spend, 0), - w.broadcast_enabled, w.talk_to_user_enabled + w.broadcast_enabled, w.talk_to_user_enabled, + COALESCE(w.compute, '{}'::jsonb) FROM workspaces w LEFT JOIN canvas_layouts cl ON cl.workspace_id = w.id WHERE w.status != 'removed' @@ -813,7 +842,8 @@ func (h *WorkspaceHandler) Get(c *gin.Context) { COALESCE(w.workspace_dir, ''), COALESCE(cl.x, 0), COALESCE(cl.y, 0), COALESCE(cl.collapsed, false), w.budget_limit, COALESCE(w.monthly_spend, 0), - w.broadcast_enabled, w.talk_to_user_enabled + w.broadcast_enabled, w.talk_to_user_enabled, + COALESCE(w.compute, '{}'::jsonb) FROM workspaces w LEFT JOIN canvas_layouts cl ON cl.workspace_id = w.id WHERE w.id = $1 diff --git a/workspace-server/internal/handlers/workspace_budget_test.go b/workspace-server/internal/handlers/workspace_budget_test.go index 4652e2932..67fd77237 100644 --- a/workspace-server/internal/handlers/workspace_budget_test.go +++ b/workspace-server/internal/handlers/workspace_budget_test.go @@ -33,7 +33,7 @@ var wsColumns = []string{ "parent_id", "active_tasks", "max_concurrent_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } // ==================== GET — financial fields stripped from open endpoint ==================== @@ -56,7 +56,8 @@ func TestWorkspaceBudget_Get_NilLimit(t *testing.T) { nil, // budget_limit NULL 0, // monthly_spend 0 false, // broadcast_enabled - true)) // talk_to_user_enabled + true, // talk_to_user_enabled + []byte(`{}`))) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -100,7 +101,8 @@ func TestWorkspaceBudget_Get_WithLimit(t *testing.T) { 0.0, 0.0, false, int64(500), // budget_limit = $5.00 in DB int64(123), // monthly_spend = $1.23 in DB - false, true)) // broadcast_enabled, talk_to_user_enabled + false, true, // broadcast_enabled, talk_to_user_enabled + []byte(`{}`))) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -145,18 +147,18 @@ func TestWorkspaceBudget_Create_WithLimit(t *testing.T) { mock.ExpectBegin() mock.ExpectExec("INSERT INTO workspaces"). WithArgs( - sqlmock.AnyArg(), // id - "Budgeted Agent", // name - nil, // role - 3, // tier (default, workspace.go create-handler) - "langgraph", // runtime - sqlmock.AnyArg(), // awareness_namespace - (*string)(nil), // parent_id - nil, // workspace_dir - "none", // workspace_access - &budgetVal, // budget_limit ($10) + sqlmock.AnyArg(), // id + "Budgeted Agent", // name + nil, // role + 3, // tier (default, workspace.go create-handler) + "langgraph", // runtime + sqlmock.AnyArg(), // awareness_namespace + (*string)(nil), // parent_id + nil, // workspace_dir + "none", // workspace_access + &budgetVal, // budget_limit ($10) models.DefaultMaxConcurrentTasks, // max_concurrent_tasks default - "push", // delivery_mode default (#2339) + "push", // delivery_mode default (#2339) ). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() diff --git a/workspace-server/internal/handlers/workspace_compute.go b/workspace-server/internal/handlers/workspace_compute.go new file mode 100644 index 000000000..58522f955 --- /dev/null +++ b/workspace-server/internal/handlers/workspace_compute.go @@ -0,0 +1,128 @@ +package handlers + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "log" + + "github.com/Molecule-AI/molecule-monorepo/platform/internal/db" + "github.com/Molecule-AI/molecule-monorepo/platform/internal/models" +) + +const ( + workspaceComputeDiskFloorGB = 30 + workspaceComputeDiskCeilingGB = 500 +) + +var workspaceComputeInstanceAllowlist = map[string]struct{}{ + "t3.medium": {}, + "t3.large": {}, + "t3.xlarge": {}, + "t3.2xlarge": {}, + "m6i.large": {}, + "m6i.xlarge": {}, + "c6i.xlarge": {}, +} + +func validateWorkspaceCompute(compute models.WorkspaceCompute) error { + if compute.InstanceType != "" { + if _, ok := workspaceComputeInstanceAllowlist[compute.InstanceType]; !ok { + return fmt.Errorf("unsupported compute.instance_type") + } + } + if compute.Volume.RootGB != 0 { + if compute.Volume.RootGB < workspaceComputeDiskFloorGB || compute.Volume.RootGB > workspaceComputeDiskCeilingGB { + return fmt.Errorf("compute.volume.root_gb must be between %d and %d", workspaceComputeDiskFloorGB, workspaceComputeDiskCeilingGB) + } + } + switch compute.Display.Mode { + case "", "none", "desktop-control", "gpu-desktop-control": + default: + return fmt.Errorf("unsupported compute.display.mode") + } + switch compute.Display.Protocol { + case "", "dcv": + default: + return fmt.Errorf("unsupported compute.display.protocol") + } + if compute.Display.Width < 0 || compute.Display.Height < 0 { + return fmt.Errorf("compute.display width/height must be non-negative") + } + return nil +} + +func workspaceComputeIsZero(compute models.WorkspaceCompute) bool { + return compute.InstanceType == "" && + compute.Volume.RootGB == 0 && + compute.Display.Mode == "" && + compute.Display.Width == 0 && + compute.Display.Height == 0 && + compute.Display.Protocol == "" +} + +func workspaceComputeJSON(compute models.WorkspaceCompute) (string, error) { + if workspaceComputeIsZero(compute) { + return "{}", nil + } + out := map[string]interface{}{} + if compute.InstanceType != "" { + out["instance_type"] = compute.InstanceType + } + if compute.Volume.RootGB != 0 { + out["volume"] = map[string]interface{}{"root_gb": compute.Volume.RootGB} + } + display := map[string]interface{}{} + if compute.Display.Mode != "" { + display["mode"] = compute.Display.Mode + } + if compute.Display.Width != 0 { + display["width"] = compute.Display.Width + } + if compute.Display.Height != 0 { + display["height"] = compute.Display.Height + } + if compute.Display.Protocol != "" { + display["protocol"] = compute.Display.Protocol + } + if len(display) > 0 { + out["display"] = display + } + b, err := json.Marshal(out) + if err != nil { + return "", err + } + return string(b), nil +} + +func withStoredCompute(ctx context.Context, workspaceID string, payload models.CreateWorkspacePayload) models.CreateWorkspacePayload { + if !workspaceComputeIsZero(payload.Compute) || db.DB == nil { + return payload + } + var raw string + err := db.DB.QueryRowContext(ctx, + `SELECT COALESCE(compute, '{}'::jsonb) FROM workspaces WHERE id = $1`, + workspaceID, + ).Scan(&raw) + if err != nil { + if err != sql.ErrNoRows { + log.Printf("withStoredCompute: load compute for %s failed: %v", workspaceID, err) + } + return payload + } + if raw == "" || raw == "{}" { + return payload + } + var compute models.WorkspaceCompute + if err := json.Unmarshal([]byte(raw), &compute); err != nil { + log.Printf("withStoredCompute: invalid compute JSON for %s: %v", workspaceID, err) + return payload + } + if err := validateWorkspaceCompute(compute); err != nil { + log.Printf("withStoredCompute: stored compute for %s failed validation: %v", workspaceID, err) + return payload + } + payload.Compute = compute + return payload +} diff --git a/workspace-server/internal/handlers/workspace_compute_test.go b/workspace-server/internal/handlers/workspace_compute_test.go new file mode 100644 index 000000000..f82c4262a --- /dev/null +++ b/workspace-server/internal/handlers/workspace_compute_test.go @@ -0,0 +1,175 @@ +package handlers + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/Molecule-AI/molecule-monorepo/platform/internal/models" + "github.com/gin-gonic/gin" +) + +func TestValidateWorkspaceCompute_AcceptsPhase1SizingAndDisplayNone(t *testing.T) { + compute := models.WorkspaceCompute{ + InstanceType: "m6i.xlarge", + Volume: models.WorkspaceComputeVolume{RootGB: 100}, + Display: models.WorkspaceComputeDisplay{Mode: "none"}, + } + + if err := validateWorkspaceCompute(compute); err != nil { + t.Fatalf("validateWorkspaceCompute returned error for valid compute: %v", err) + } +} + +func TestValidateWorkspaceCompute_RejectsUnknownInstanceType(t *testing.T) { + compute := models.WorkspaceCompute{InstanceType: "p4d.24xlarge"} + + if err := validateWorkspaceCompute(compute); err == nil { + t.Fatal("validateWorkspaceCompute accepted unsupported instance type") + } +} + +func TestValidateWorkspaceCompute_RejectsOutOfRangeRootVolume(t *testing.T) { + for _, rootGB := range []int{29, 501} { + compute := models.WorkspaceCompute{Volume: models.WorkspaceComputeVolume{RootGB: rootGB}} + if err := validateWorkspaceCompute(compute); err == nil { + t.Fatalf("validateWorkspaceCompute accepted root_gb=%d", rootGB) + } + } +} + +func TestWorkspaceComputeJSON_OmitsEmptyNestedSections(t *testing.T) { + got, err := workspaceComputeJSON(models.WorkspaceCompute{ + InstanceType: "m6i.xlarge", + Volume: models.WorkspaceComputeVolume{RootGB: 100}, + }) + if err != nil { + t.Fatalf("workspaceComputeJSON returned error: %v", err) + } + + if strings.Contains(got, `"display"`) { + t.Fatalf("workspaceComputeJSON included empty display section: %s", got) + } + if got != `{"instance_type":"m6i.xlarge","volume":{"root_gb":100}}` { + t.Fatalf("workspaceComputeJSON = %s", got) + } +} + +func TestWorkspaceCreate_WithCompute_PersistsComputeJSON(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + broadcaster := newTestBroadcaster() + handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) + + mock.ExpectBegin() + mock.ExpectExec("INSERT INTO workspaces"). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectExec(`UPDATE workspaces SET compute = \$2::jsonb`). + WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectCommit() + mock.ExpectExec("INSERT INTO canvas_layouts"). + WillReturnResult(sqlmock.NewResult(0, 1)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + body := `{ + "name":"Sized Agent", + "external":true, + "runtime":"external", + "compute":{ + "instance_type":"m6i.xlarge", + "volume":{"root_gb":100}, + "display":{"mode":"none"} + } + }` + c.Request = httptest.NewRequest("POST", "/workspaces", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + + handler.Create(c) + + if w.Code != http.StatusCreated { + t.Fatalf("expected status 201, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceCreate_WithInvalidCompute_ReturnsBadRequest(t *testing.T) { + setupTestDB(t) + setupTestRedis(t) + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + body := `{ + "name":"Oversized Agent", + "compute":{"instance_type":"p4d.24xlarge"} + }` + c.Request = httptest.NewRequest("POST", "/workspaces", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + + handler.Create(c) + + if w.Code != http.StatusBadRequest { + t.Fatalf("expected status 400, got %d: %s", w.Code, w.Body.String()) + } +} + +func TestBuildProvisionerConfig_CopiesComputeSizingFromPayload(t *testing.T) { + mock := setupTestDB(t) + mock.ExpectQuery(`SELECT COALESCE\(workspace_dir`). + WithArgs("ws-compute"). + WillReturnRows(sqlmock.NewRows([]string{"workspace_dir", "workspace_access"}).AddRow("", "none")) + + handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + cfg := handler.buildProvisionerConfig( + context.Background(), + "ws-compute", + "", + nil, + models.CreateWorkspacePayload{ + Tier: 4, + Runtime: "claude-code", + Compute: models.WorkspaceCompute{ + InstanceType: "m6i.xlarge", + Volume: models.WorkspaceComputeVolume{RootGB: 100}, + }, + }, + nil, + t.TempDir(), + "workspace:ws-compute", + ) + + if cfg.InstanceType != "m6i.xlarge" { + t.Errorf("cfg.InstanceType = %q, want m6i.xlarge", cfg.InstanceType) + } + if cfg.DiskGB != 100 { + t.Errorf("cfg.DiskGB = %d, want 100", cfg.DiskGB) + } +} + +func TestWithStoredCompute_LoadsComputeForRestartPayloads(t *testing.T) { + mock := setupTestDB(t) + mock.ExpectQuery(`SELECT COALESCE\(compute, '\{\}'::jsonb\) FROM workspaces WHERE id = \$1`). + WithArgs("ws-restart-compute"). + WillReturnRows(sqlmock.NewRows([]string{"compute"}).AddRow(`{"instance_type":"m6i.xlarge","volume":{"root_gb":100}}`)) + + payload := models.CreateWorkspacePayload{Name: "Restart Me", Tier: 4, Runtime: "claude-code"} + got := withStoredCompute(context.Background(), "ws-restart-compute", payload) + + if got.Compute.InstanceType != "m6i.xlarge" { + t.Errorf("stored compute instance_type = %q, want m6i.xlarge", got.Compute.InstanceType) + } + if got.Compute.Volume.RootGB != 100 { + t.Errorf("stored compute root_gb = %d, want 100", got.Compute.Volume.RootGB) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} diff --git a/workspace-server/internal/handlers/workspace_provision.go b/workspace-server/internal/handlers/workspace_provision.go index 11da5f448..2ba89a740 100644 --- a/workspace-server/internal/handlers/workspace_provision.go +++ b/workspace-server/internal/handlers/workspace_provision.go @@ -296,6 +296,8 @@ func (h *WorkspaceHandler) buildProvisionerConfig( WorkspaceAccess: workspaceAccess, Tier: payload.Tier, Runtime: payload.Runtime, + InstanceType: payload.Compute.InstanceType, + DiskGB: int32(payload.Compute.Volume.RootGB), EnvVars: envVars, PlatformURL: h.platformURL, AwarenessURL: os.Getenv("AWARENESS_URL"), @@ -1009,4 +1011,3 @@ func (h *WorkspaceHandler) provisionWorkspaceCP(workspaceID, templatePath string log.Printf("CPProvisioner: workspace %s started as machine %s via control plane", workspaceID, machineID) } - diff --git a/workspace-server/internal/handlers/workspace_restart.go b/workspace-server/internal/handlers/workspace_restart.go index 6eb490913..6ac3dc506 100644 --- a/workspace-server/internal/handlers/workspace_restart.go +++ b/workspace-server/internal/handlers/workspace_restart.go @@ -164,7 +164,7 @@ func (h *WorkspaceHandler) maybeRestartAfterFileWrite(workspaceID string) { // isRestarting reports whether a restart cycle is currently in flight for // the workspace. Callers that have their own "container looks dead" probe // MUST consult this before triggering a restart, because during the -// 20-30s EC2-pending window the workspace's url='' and IsRunning()=false +// 20-30s EC2-pending window the workspace's url=” and IsRunning()=false // looks identical to a dead container — and any restart-triggering probe // (maybeMarkContainerDead from canvas /delegations poll, or the trailing // restart-context probe at the end of runRestartCycle) will set @@ -337,7 +337,7 @@ func (h *WorkspaceHandler) Restart(c *gin.Context) { } var configFiles map[string][]byte - payload := models.CreateWorkspacePayload{Name: wsName, Tier: tier, Runtime: containerRuntime} + payload := withStoredCompute(ctx, id, models.CreateWorkspacePayload{Name: wsName, Tier: tier, Runtime: containerRuntime}) log.Printf("Restart: workspace %s (%s) runtime=%q", wsName, id, containerRuntime) // #12: ?reset=true (or body.Reset) discards the claude-sessions volume @@ -791,7 +791,7 @@ func (h *WorkspaceHandler) runRestartCycle(workspaceID string) { }) // Runtime from DB — no more config file parsing - payload := models.CreateWorkspacePayload{Name: wsName, Tier: tier, Runtime: dbRuntime} + payload := withStoredCompute(ctx, workspaceID, models.CreateWorkspacePayload{Name: wsName, Tier: tier, Runtime: dbRuntime}) // Snapshot restart-context data before the new session overwrites // last_heartbeat_at. Issue #19 Layer 1. @@ -948,7 +948,7 @@ func (h *WorkspaceHandler) Resume(c *gin.Context) { h.broadcaster.RecordAndBroadcast(ctx, string(events.EventWorkspaceProvisioning), ws.id, map[string]interface{}{ "name": ws.name, "tier": ws.tier, "runtime": ws.runtime, }) - payload := models.CreateWorkspacePayload{Name: ws.name, Tier: ws.tier, Runtime: ws.runtime} + payload := withStoredCompute(ctx, ws.id, models.CreateWorkspacePayload{Name: ws.name, Tier: ws.tier, Runtime: ws.runtime}) // Resume is provision-only (workspace is paused, no live container // to stop). provisionWorkspaceAuto handles backend routing and the // no-backend mark-failed fallback identically to Create. Pre- diff --git a/workspace-server/internal/handlers/workspace_test.go b/workspace-server/internal/handlers/workspace_test.go index 7f329da2e..e8deeb475 100644 --- a/workspace-server/internal/handlers/workspace_test.go +++ b/workspace-server/internal/handlers/workspace_test.go @@ -29,7 +29,7 @@ func TestWorkspaceGet_Success(t *testing.T) { "parent_id", "active_tasks", "max_concurrent_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } mock.ExpectQuery("SELECT w.id, w.name"). WithArgs("cccccccc-0001-0000-0000-000000000000"). @@ -37,7 +37,7 @@ func TestWorkspaceGet_Success(t *testing.T) { AddRow("cccccccc-0001-0000-0000-000000000000", "My Agent", "worker", 1, "online", []byte(`{"name":"test"}`), "http://localhost:8001", nil, 2, 1, 0.05, "", 3600, "working", "langgraph", "", 10.0, 20.0, false, - nil, 0, false, true)) + nil, 0, false, true, []byte(`{}`))) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -119,7 +119,7 @@ func TestWorkspaceGet_RemovedReturns410(t *testing.T) { "parent_id", "active_tasks", "max_concurrent_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } mock.ExpectQuery("SELECT w.id, w.name"). WithArgs(id). @@ -127,7 +127,7 @@ func TestWorkspaceGet_RemovedReturns410(t *testing.T) { AddRow(id, "Old Agent", "worker", 1, string(models.StatusRemoved), []byte(`null`), "", nil, 0, 1, 0.0, "", 0, "", "langgraph", "", 0.0, 0.0, false, - nil, 0, false, true)) + nil, 0, false, true, []byte(`{}`))) mock.ExpectQuery(`SELECT updated_at FROM workspaces`). WithArgs(id). WillReturnRows(sqlmock.NewRows([]string{"updated_at"}).AddRow(removedAt)) @@ -183,7 +183,7 @@ func TestWorkspaceGet_RemovedReturns410WithNullRemovedAtOnTimestampFetchFailure( "parent_id", "active_tasks", "max_concurrent_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } mock.ExpectQuery("SELECT w.id, w.name"). WithArgs(id). @@ -191,7 +191,7 @@ func TestWorkspaceGet_RemovedReturns410WithNullRemovedAtOnTimestampFetchFailure( AddRow(id, "Vanished", "worker", 1, string(models.StatusRemoved), []byte(`null`), "", nil, 0, 1, 0.0, "", 0, "", "langgraph", "", 0.0, 0.0, false, - nil, 0, false, true)) + nil, 0, false, true, []byte(`{}`))) // Simulate the row vanishing between the two queries. mock.ExpectQuery(`SELECT updated_at FROM workspaces`). WithArgs(id). @@ -246,7 +246,7 @@ func TestWorkspaceGet_RemovedWithIncludeQueryReturns200(t *testing.T) { "parent_id", "active_tasks", "max_concurrent_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } mock.ExpectQuery("SELECT w.id, w.name"). WithArgs(id). @@ -254,7 +254,7 @@ func TestWorkspaceGet_RemovedWithIncludeQueryReturns200(t *testing.T) { AddRow(id, "Audit Agent", "worker", 1, string(models.StatusRemoved), []byte(`null`), "", nil, 0, 1, 0.0, "", 0, "", "langgraph", "", 0.0, 0.0, false, - nil, 0, false, true)) + nil, 0, false, true, []byte(`{}`))) // last_outbound_at follow-up query (existing path) mock.ExpectQuery(`SELECT last_outbound_at FROM workspaces`). WithArgs(id). @@ -718,7 +718,7 @@ func TestWorkspaceList_Empty(t *testing.T) { "parent_id", "active_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", })) w := httptest.NewRecorder() @@ -1422,7 +1422,7 @@ func TestWorkspaceGet_FinancialFieldsStripped(t *testing.T) { "parent_id", "active_tasks", "max_concurrent_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } // Populate with non-zero financial values to confirm they are stripped. mock.ExpectQuery("SELECT w.id, w.name"). @@ -1431,7 +1431,7 @@ func TestWorkspaceGet_FinancialFieldsStripped(t *testing.T) { AddRow("cccccccc-0010-0000-0000-000000000000", "Finance Test", "worker", 1, "online", []byte(`{}`), "http://localhost:9001", nil, 0, 1, 0.0, "", 0, "", "langgraph", "", 0.0, 0.0, false, - int64(50000), int64(12500), false, true)) // budget_limit=500 USD, spend=125 USD + int64(50000), int64(12500), false, true, []byte(`{}`))) // budget_limit=500 USD, spend=125 USD w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -1479,7 +1479,7 @@ func TestWorkspaceGet_SensitiveFieldsStripped(t *testing.T) { "parent_id", "active_tasks", "max_concurrent_tasks", "last_error_rate", "last_sample_error", "uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed", "budget_limit", "monthly_spend", - "broadcast_enabled", "talk_to_user_enabled", + "broadcast_enabled", "talk_to_user_enabled", "compute", } mock.ExpectQuery("SELECT w.id, w.name"). WithArgs("cccccccc-0955-0000-0000-000000000000"). @@ -1492,7 +1492,7 @@ func TestWorkspaceGet_SensitiveFieldsStripped(t *testing.T) { "langgraph", "/home/user/secret-projects/client-work", 0.0, 0.0, false, - nil, 0, false, true)) + nil, 0, false, true, []byte(`{}`))) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) diff --git a/workspace-server/internal/models/workspace.go b/workspace-server/internal/models/workspace.go index 9139fc5b9..ecdc7a023 100644 --- a/workspace-server/internal/models/workspace.go +++ b/workspace-server/internal/models/workspace.go @@ -35,16 +35,16 @@ type Workspace struct { // DeliveryMode: "push" (synchronous to URL — default) or "poll" (logged // to activity_logs, agent reads via GET /activity?since_id=). See // migration 045 + RFC #2339. - DeliveryMode string `json:"delivery_mode" db:"delivery_mode"` + DeliveryMode string `json:"delivery_mode" db:"delivery_mode"` // BroadcastEnabled: when true the workspace may call POST /broadcast to // deliver a message to all non-removed agent workspaces in the org. // Default false — only privileged orchestrators should hold this ability. - BroadcastEnabled bool `json:"broadcast_enabled" db:"broadcast_enabled"` + BroadcastEnabled bool `json:"broadcast_enabled" db:"broadcast_enabled"` // TalkToUserEnabled: when false the workspace's send_message_to_user calls // and POST /notify requests are rejected with HTTP 403 so the agent is // forced to route updates through a parent workspace. Default true // (preserves existing behaviour for all workspaces). - TalkToUserEnabled bool `json:"talk_to_user_enabled" db:"talk_to_user_enabled"` + TalkToUserEnabled bool `json:"talk_to_user_enabled" db:"talk_to_user_enabled"` // Canvas layout fields (from JOIN) X float64 `json:"x"` Y float64 `json:"y"` @@ -71,12 +71,12 @@ type RegisterPayload struct { // enforces the conditional requirement based on the resolved // delivery mode (payload value, falling back to the row's existing // value, falling back to "push"). - URL string `json:"url"` - AgentCard json.RawMessage `json:"agent_card" binding:"required"` + URL string `json:"url"` + AgentCard json.RawMessage `json:"agent_card" binding:"required"` // DeliveryMode is optional. Empty string means "keep the existing // value on the workspace row, or default to push for new rows". // When set, must be one of DeliveryModePush / DeliveryModePoll. - DeliveryMode string `json:"delivery_mode,omitempty"` + DeliveryMode string `json:"delivery_mode,omitempty"` } type HeartbeatPayload struct { @@ -154,19 +154,36 @@ type MemorySeed struct { Scope string `json:"scope" yaml:"scope"` // LOCAL, TEAM, GLOBAL } +type WorkspaceComputeVolume struct { + RootGB int `json:"root_gb,omitempty"` +} + +type WorkspaceComputeDisplay struct { + Mode string `json:"mode,omitempty"` + Width int `json:"width,omitempty"` + Height int `json:"height,omitempty"` + Protocol string `json:"protocol,omitempty"` +} + +type WorkspaceCompute struct { + InstanceType string `json:"instance_type,omitempty"` + Volume WorkspaceComputeVolume `json:"volume,omitempty"` + Display WorkspaceComputeDisplay `json:"display,omitempty"` +} + type CreateWorkspacePayload struct { - Name string `json:"name" binding:"required"` - Role string `json:"role"` - Template string `json:"template"` // workspace-configs-templates folder name - Tier int `json:"tier"` - Model string `json:"model"` - Runtime string `json:"runtime"` // "langgraph" (default), "claude-code", etc. - External bool `json:"external"` // true = no Docker container, just a registered URL - URL string `json:"url"` // for external workspaces: the A2A endpoint URL (push mode only — omit for poll) + Name string `json:"name" binding:"required"` + Role string `json:"role"` + Template string `json:"template"` // workspace-configs-templates folder name + Tier int `json:"tier"` + Model string `json:"model"` + Runtime string `json:"runtime"` // "langgraph" (default), "claude-code", etc. + External bool `json:"external"` // true = no Docker container, just a registered URL + URL string `json:"url"` // for external workspaces: the A2A endpoint URL (push mode only — omit for poll) // DeliveryMode: "push" (default) sends inbound A2A to URL synchronously; // "poll" records inbound to activity_logs for the agent to consume via // GET /activity?since_id=. Poll mode does not require a URL. See #2339. - DeliveryMode string `json:"delivery_mode,omitempty"` + DeliveryMode string `json:"delivery_mode,omitempty"` WorkspaceDir string `json:"workspace_dir"` // host path to mount as /workspace (empty = isolated volume) WorkspaceAccess string `json:"workspace_access"` // "none" (default), "read_only", or "read_write" — see #65 ParentID *string `json:"parent_id"` @@ -180,7 +197,11 @@ type CreateWorkspacePayload struct { // MaxConcurrentTasks caps parallel A2A + cron dispatch. 0 means use // DefaultMaxConcurrentTasks. Leaders typically set 3. MaxConcurrentTasks int `json:"max_concurrent_tasks"` - Canvas struct { + // Compute is the product-facing per-workspace EC2 shape/display + // contract. Phase 1 uses instance_type + volume.root_gb and persists + // display for future desktop-control workspaces. + Compute WorkspaceCompute `json:"compute,omitempty"` + Canvas struct { X float64 `json:"x"` Y float64 `json:"y"` } `json:"canvas"` diff --git a/workspace-server/internal/provisioner/cp_provisioner.go b/workspace-server/internal/provisioner/cp_provisioner.go index 8f6f0c557..269816909 100644 --- a/workspace-server/internal/provisioner/cp_provisioner.go +++ b/workspace-server/internal/provisioner/cp_provisioner.go @@ -152,12 +152,14 @@ func (p *CPProvisioner) adminAuthHeaders(req *http.Request) { } type cpProvisionRequest struct { - OrgID string `json:"org_id"` - WorkspaceID string `json:"workspace_id"` - Runtime string `json:"runtime"` - Tier int `json:"tier"` - PlatformURL string `json:"platform_url"` - Env map[string]string `json:"env"` + OrgID string `json:"org_id"` + WorkspaceID string `json:"workspace_id"` + Runtime string `json:"runtime"` + Tier int `json:"tier"` + InstanceType string `json:"instance_type,omitempty"` + DiskGB int32 `json:"disk_gb,omitempty"` + PlatformURL string `json:"platform_url"` + Env map[string]string `json:"env"` // ConfigFiles are template + generated config files to write into the // EC2 instance's /configs directory. OFFSEC-010: collected by // collectCPConfigFiles which rejects symlinks and non-regular files @@ -206,13 +208,15 @@ func (p *CPProvisioner) Start(ctx context.Context, cfg WorkspaceConfig) (string, } req := cpProvisionRequest{ - OrgID: p.orgID, - WorkspaceID: cfg.WorkspaceID, - Runtime: cfg.Runtime, - Tier: cfg.Tier, - PlatformURL: cfg.PlatformURL, - Env: env, - ConfigFiles: configFiles, + OrgID: p.orgID, + WorkspaceID: cfg.WorkspaceID, + Runtime: cfg.Runtime, + Tier: cfg.Tier, + InstanceType: cfg.InstanceType, + DiskGB: cfg.DiskGB, + PlatformURL: cfg.PlatformURL, + Env: env, + ConfigFiles: configFiles, } body, err := json.Marshal(req) diff --git a/workspace-server/internal/provisioner/cp_provisioner_test.go b/workspace-server/internal/provisioner/cp_provisioner_test.go index 6f1ea07e8..44e27660a 100644 --- a/workspace-server/internal/provisioner/cp_provisioner_test.go +++ b/workspace-server/internal/provisioner/cp_provisioner_test.go @@ -191,6 +191,12 @@ func TestStart_HappyPath(t *testing.T) { if body.WorkspaceID != "ws-1" || body.Runtime != "python" { t.Errorf("body mismatch: %+v", body) } + if body.InstanceType != "m6i.xlarge" { + t.Errorf("instance_type = %q, want m6i.xlarge", body.InstanceType) + } + if body.DiskGB != 100 { + t.Errorf("disk_gb = %d, want 100", body.DiskGB) + } w.WriteHeader(http.StatusCreated) _, _ = io.WriteString(w, `{"instance_id":"i-abc123","state":"pending"}`) })) @@ -205,6 +211,7 @@ func TestStart_HappyPath(t *testing.T) { id, err := p.Start(context.Background(), WorkspaceConfig{ WorkspaceID: "ws-1", Runtime: "python", Tier: 1, PlatformURL: "http://tenant", + InstanceType: "m6i.xlarge", DiskGB: 100, }) if err != nil { t.Fatalf("Start: %v", err) @@ -362,7 +369,7 @@ func TestStart_CollectsConfigFiles(t *testing.T) { p := &CPProvisioner{baseURL: srv.URL, orgID: "org-1", httpClient: srv.Client()} _, err := p.Start(context.Background(), WorkspaceConfig{ WorkspaceID: "ws-1", - Runtime: "python", + Runtime: "python", Tier: 1, PlatformURL: "http://tenant", TemplatePath: tmpl, @@ -424,7 +431,7 @@ func TestStart_SymlinkTemplatePathError(t *testing.T) { p := &CPProvisioner{baseURL: "http://unused", orgID: "org-1", httpClient: &http.Client{Timeout: time.Second}} _, err := p.Start(context.Background(), WorkspaceConfig{ WorkspaceID: "ws-1", - Runtime: "python", + Runtime: "python", TemplatePath: symlink, // symlink root → OFFSEC-010 guard should fire }) if err == nil { diff --git a/workspace-server/internal/provisioner/provisioner.go b/workspace-server/internal/provisioner/provisioner.go index 164c951bf..0dbfee309 100644 --- a/workspace-server/internal/provisioner/provisioner.go +++ b/workspace-server/internal/provisioner/provisioner.go @@ -98,6 +98,8 @@ type WorkspaceConfig struct { WorkspacePath string // Host path to bind-mount as /workspace (if empty, uses Docker named volume) Tier int Runtime string // "langgraph" (default) or "claude-code", "codex", "ollama", "custom" + InstanceType string // Optional CP EC2 instance type override (SaaS only) + DiskGB int32 // Optional CP root volume size override in GiB (SaaS only) EnvVars map[string]string // Additional env vars (API keys, etc.) PlatformURL string AwarenessURL string @@ -1605,4 +1607,3 @@ func parseOCIPlatform(s string) *ocispec.Platform { } return &ocispec.Platform{OS: parts[0], Architecture: parts[1]} } - diff --git a/workspace-server/migrations/050_workspace_compute.down.sql b/workspace-server/migrations/050_workspace_compute.down.sql new file mode 100644 index 000000000..f53d248d8 --- /dev/null +++ b/workspace-server/migrations/050_workspace_compute.down.sql @@ -0,0 +1 @@ +ALTER TABLE workspaces DROP COLUMN IF EXISTS compute; diff --git a/workspace-server/migrations/050_workspace_compute.up.sql b/workspace-server/migrations/050_workspace_compute.up.sql new file mode 100644 index 000000000..97f91afe6 --- /dev/null +++ b/workspace-server/migrations/050_workspace_compute.up.sql @@ -0,0 +1,2 @@ +ALTER TABLE workspaces + ADD COLUMN IF NOT EXISTS compute JSONB NOT NULL DEFAULT '{}'::jsonb;