direct: simplify Engine API: SaveState(id, state) replaces SetID+SaveState

Collapse the two-step engine.SetID(id) / engine.SaveState(state) pattern into a
single engine.SaveState(id, state) call. The Engine still tracks the id internally
and panics if a subsequent call passes a different id (guards against bugs).

Also extract the polling loop from vector_search_index.DoCreate into a private
waitForIndexReady helper so that the unchanged DoDelete stays between the two
functions and does not appear in the diff.

Co-authored-by: Denis Bilenko

Author: denik

bundle/direct/apply.go

@@ -129,8 +129,6 @@ func (d *DeploymentUnit) Update(ctx context.Context, db *dstate.DeploymentState,
 	engine := dresources.NewEngine(d.Adapter.StateType(), func(_ string, x any) error {
 		return db.SaveState(d.ResourceKey, id, x, d.DependsOn)
 	})
-	engine.SetID(id)
-
 	remoteState, err := retryOnTransient(ctx, func() (any, error) {
 		return d.Adapter.DoUpdate(ctx, engine, id, newState, planEntry)
 	})

bundle/direct/dresources/adapter.go

@@ -53,7 +53,7 @@ type IResource interface {
 
 	// DoCreate creates a new resource from the newState. Returns id of the resource and optionally remote state.
 	// If remote state is available as part of the operation, return it; otherwise return nil.
-	// Call engine.SetID then engine.SaveState to persist intermediate state before long-running waits.
+	// Call engine.SaveState(id, state) to persist intermediate state before long-running waits.
 	// Example: func (r *ResourceVolume) DoCreate(ctx context.Context, _ *Engine, newState *catalog.CreateVolumeRequestContent) (string, *catalog.VolumeInfo, error)
 	DoCreate(ctx context.Context, engine *Engine, newState any) (id string, remoteState any, e error)
 

bundle/direct/dresources/app.go

@@ -156,8 +156,8 @@ func (r *ResourceApp) DoCreate(ctx context.Context, engine *Engine, config *AppS
 
 	// Save state as soon as the app exists so it is not orphaned if the wait or
 	// lifecycle management is interrupted.
-	engine.SetID(app.Name)
-	if err := engine.SaveState(config); err != nil {
+
+	if err := engine.SaveState(app.Name, config); err != nil {
 		return "", nil, err
 	}
 

bundle/direct/dresources/cluster.go

@@ -156,8 +156,7 @@ func (r *ResourceCluster) DoCreate(ctx context.Context, engine *Engine, config *
 
 	// Save state immediately after the cluster is created so it is not orphaned
 	// if the subsequent wait or terminate is interrupted.
-	engine.SetID(id)
-	if err := engine.SaveState(config); err != nil {
+	if err := engine.SaveState(id, config); err != nil {
 		return "", nil, err
 	}
 

bundle/direct/dresources/database_instance.go

@@ -36,8 +36,7 @@ func (d *ResourceDatabaseInstance) DoCreate(ctx context.Context, engine *Engine,
 
 	// Save state immediately after the instance is created so it is not orphaned
 	// if the subsequent wait is interrupted.
-	engine.SetID(id)
-	if err := engine.SaveState(config); err != nil {
+	if err := engine.SaveState(id, config); err != nil {
 		return "", nil, err
 	}
 

bundle/direct/dresources/engine.go

@@ -1,7 +1,6 @@
 package dresources
 
 import (
-	"errors"
 	"fmt"
 	"reflect"
 )
@@ -26,18 +25,14 @@ func NewNopEngine(stateType reflect.Type) *Engine {
 	return NewEngine(stateType, func(_ string, _ any) error { return nil })
 }
 
-// SetID sets the resource id for subsequent SaveState calls.
-// Must be called before SaveState during DoCreate; for DoUpdate the Engine is
-// pre-configured with the existing id.
-func (e *Engine) SetID(id string) {
-	e.id = id
-}
-
-// SaveState saves the resource state. x must be of the same pointer-to-struct
-// type as the resource's state type. Returns an error if SetID was not called.
-func (e *Engine) SaveState(x any) error {
+// SaveState saves the resource state. id must be the resource's identifier; on
+// the first call it is recorded, and subsequent calls panic if a different id is
+// passed. x must be a pointer to the same struct type as the resource's state.
+func (e *Engine) SaveState(id string, x any) error {
 	if e.id == "" {
-		return errors.New("SaveState: id not set, call SetID first")
+		e.id = id
+	} else if e.id != id {
+		panic(fmt.Sprintf("SaveState: id mismatch: expected %q, got %q", e.id, id))
 	}
 	xt := reflect.TypeOf(x)
 	if xt != e.stateType {

bundle/direct/dresources/model_serving_endpoint.go

@@ -137,8 +137,7 @@ func (r *ResourceModelServingEndpoint) DoCreate(ctx context.Context, engine *Eng
 
 	// Save state immediately after the endpoint is created so it is not orphaned
 	// if the subsequent wait is interrupted.
-	engine.SetID(id)
-	if err := engine.SaveState(config); err != nil {
+	if err := engine.SaveState(id, config); err != nil {
 		return "", nil, err
 	}
 

bundle/direct/dresources/vector_search_endpoint.go

@@ -74,8 +74,7 @@ func (r *ResourceVectorSearchEndpoint) DoCreate(ctx context.Context, engine *Eng
 
 	// Save state immediately after the endpoint is created so it is not orphaned
 	// if the subsequent wait is interrupted.
-	engine.SetID(id)
-	if err := engine.SaveState(config); err != nil {
+	if err := engine.SaveState(id, config); err != nil {
 		return "", nil, err
 	}
 

bundle/direct/dresources/vector_search_index.go

@@ -138,15 +138,34 @@ func (r *ResourceVectorSearchIndex) DoCreate(ctx context.Context, engine *Engine
 
 	// Save state immediately after the index is created (endpoint UUID now set) so it
 	// is not orphaned if the subsequent provisioning wait is interrupted.
-	engine.SetID(config.Name)
-	if err := engine.SaveState(config); err != nil {
+	if err := engine.SaveState(config.Name, config); err != nil {
 		return "", nil, err
 	}
 
-	// CreateIndex returns immediately; poll until the embedding pipeline is ready so
-	// dependent resources and the next plan see a usable index.
+	remote, err := r.waitForIndexReady(ctx, config.Name, endpointUuid)
+	if err != nil {
+		return "", nil, err
+	}
+	return config.Name, remote, nil
+}
+
+// No DoUpdate: vector search indexes have no update API. All SDK fields are
+// declared in resources.yml under recreate_on_changes or ignore_remote_changes.
+// If a future SDK bump adds a new field that isn't classified, the framework
+// rejects the resulting Update plan at bundle_plan.go (see also the reflection
+// test in vector_search_index_test.go which catches it earlier at unit-test time).
+
+func (r *ResourceVectorSearchIndex) DoDelete(ctx context.Context, id string, _ *VectorSearchIndexState) error {
+	return r.client.VectorSearchIndexes.DeleteIndexByIndexName(ctx, id)
+}
+
+// waitForIndexReady polls GetIndex until Status.Ready=true. CreateIndex returns
+// immediately with metadata of an index whose embedding pipeline is still
+// provisioning; queries against an index that isn't ready fail. Blocking here
+// lets dependent resources (and the next plan) see a usable index.
+func (r *ResourceVectorSearchIndex) waitForIndexReady(ctx context.Context, id, endpointUuid string) (*VectorSearchIndexRemote, error) {
 	index, err := retries.Poll(ctx, createIndexTimeout, func() (*vectorsearch.VectorIndex, *retries.Err) {
-		idx, getErr := r.client.VectorSearchIndexes.GetIndexByIndexName(ctx, config.Name)
+		idx, getErr := r.client.VectorSearchIndexes.GetIndexByIndexName(ctx, id)
 		if getErr != nil {
 			return nil, retries.Halt(getErr)
 		}
@@ -160,19 +179,9 @@ func (r *ResourceVectorSearchIndex) DoCreate(ctx context.Context, engine *Engine
 		return idx, nil
 	})
 	if err != nil {
-		return "", nil, err
+		return nil, err
 	}
-	return config.Name, &VectorSearchIndexRemote{VectorIndex: *index, EndpointUuid: endpointUuid}, nil
-}
-
-// No DoUpdate: vector search indexes have no update API. All SDK fields are
-// declared in resources.yml under recreate_on_changes or ignore_remote_changes.
-// If a future SDK bump adds a new field that isn't classified, the framework
-// rejects the resulting Update plan at bundle_plan.go (see also the reflection
-// test in vector_search_index_test.go which catches it earlier at unit-test time).
-
-func (r *ResourceVectorSearchIndex) DoDelete(ctx context.Context, id string, _ *VectorSearchIndexState) error {
-	return r.client.VectorSearchIndexes.DeleteIndexByIndexName(ctx, id)
+	return &VectorSearchIndexRemote{VectorIndex: *index, EndpointUuid: endpointUuid}, nil
 }
 
 // WaitAfterDelete polls GetIndex until it returns 404. The DELETE call is