Skip to content

fix: change the model connection, readiness, and version atomically (WI-1803) #428

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stefan-ernst merged 1 commit into from
Nov 2, 2026
Merged

fix: change the model connection, readiness, and version atomically (WI-1803) #428

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 36 additions & 36 deletions internal/repository/workspace_agent_binding_repository.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -252,15 +252,17 @@ func (r *WorkspaceAgentBindingRepository) FindByActingUser(ctx context.Context,
return b, nil
}

// SetLifecycle changes only persisted lifecycle state. Runtime-affecting
// configuration methods own profile-version increments; validating a Draft
// and marking it Ready does not create another definition version.
func (r *WorkspaceAgentBindingRepository) SetLifecycle(ctx context.Context, id, workspaceID int, lifecycle models.AgentLifecycle) (int64, error) {
// SetLifecycleIfVersion changes persisted lifecycle state only when the profile
// version still matches the one that was validated. A concurrent configuration
// change bumps the version, so this refuses to mark a changed definition Ready
// using validation of an older one. Runtime-affecting configuration methods own
// profile-version increments; activating a Draft does not create another version.
func (r *WorkspaceAgentBindingRepository) SetLifecycleIfVersion(ctx context.Context, id, workspaceID int, lifecycle models.AgentLifecycle, expectedVersion int) (int64, error) {
res, err := r.db.ExecWriteContext(ctx, `
UPDATE workspace_agent_bindings
SET lifecycle = ?, updated_at = CURRENT_TIMESTAMP
WHERE id = ? AND workspace_id = ? AND lifecycle <> 'archived'
`, lifecycle, id, workspaceID)
WHERE id = ? AND workspace_id = ? AND lifecycle <> 'archived' AND profile_version = ?
`, lifecycle, id, workspaceID, expectedVersion)
if err != nil {
return 0, fmt.Errorf("set binding lifecycle: %w", err)
}
Expand Down Expand Up @@ -361,19 +363,34 @@ func (r *WorkspaceAgentBindingRepository) Restore(ctx context.Context, id, works
return n, nil
}

// UpdateInstructions rewrites a binding's custom instructions, scoped by
// workspace (WI-258).
func (r *WorkspaceAgentBindingRepository) UpdateInstructions(ctx context.Context, id, workspaceID int, instructions string) error {
_, err := r.db.ExecWriteContext(ctx, `
UPDATE workspace_agent_bindings
SET instructions = ?,
lifecycle = CASE WHEN lifecycle = 'archived' THEN lifecycle ELSE 'draft' END,
profile_version = profile_version + 1,
updated_at = CURRENT_TIMESTAMP
WHERE id = ? AND workspace_id = ?
`, instructions, id, workspaceID)
if err != nil {
return fmt.Errorf("update binding instructions: %w", err)
// UpdateAgentConfig atomically rewrites the Agent Studio overview fields —
// instructions, model connection, and runner image — and, in the same
// statement, invalidates readiness and bumps the profile version once. One
// write keeps a concurrent activation from validating the old connection and
// marking the profile Ready before the new connection is stored. A nil
// connection or runner image preserves the current value; a non-nil empty
// runner image clears the override.
func (r *WorkspaceAgentBindingRepository) UpdateAgentConfig(ctx context.Context, id, workspaceID int, instructions string, llmConnectionID *int, runnerImage *string) error {
set := []string{
"instructions = ?",
"lifecycle = CASE WHEN lifecycle = 'archived' THEN lifecycle ELSE 'draft' END",
"profile_version = profile_version + 1",
"updated_at = CURRENT_TIMESTAMP",
}
args := []any{instructions}
if llmConnectionID != nil {
set = append(set, "llm_connection_id = ?")
args = append(args, *llmConnectionID)
}
if runnerImage != nil {
set = append(set, "runner_image = ?")
args = append(args, nullStringArg(*runnerImage))
}
args = append(args, id, workspaceID)
//nolint:gosec // G201: only fixed column fragments are interpolated; every value is bound
query := "UPDATE workspace_agent_bindings SET " + strings.Join(set, ", ") + " WHERE id = ? AND workspace_id = ?"
if _, err := r.db.ExecWriteContext(ctx, query, args...); err != nil {
return fmt.Errorf("update agent config: %w", err)
}
return nil
}
Expand Down Expand Up @@ -425,23 +442,6 @@ func (r *WorkspaceAgentBindingRepository) UpdateConfig(ctx context.Context, b *m
return nil
}

// UpdateLLMConnection rewrites a binding's model connection in place, scoped by
// (id, workspace_id). It does not touch the lifecycle or profile version: the
// caller's instruction rewrite in UpdateAgentConfig owns the single Draft/version
// transition for the whole configuration change.
func (r *WorkspaceAgentBindingRepository) UpdateLLMConnection(ctx context.Context, id, workspaceID, connectionID int) error {
_, err := r.db.ExecWriteContext(ctx, `
UPDATE workspace_agent_bindings
SET llm_connection_id = ?,
updated_at = CURRENT_TIMESTAMP
WHERE id = ? AND workspace_id = ?
`, connectionID, id, workspaceID)
if err != nil {
return fmt.Errorf("update binding llm connection: %w", err)
}
return nil
}

// UpdateRunnerImage rewrites a binding's custom runner image, scoped by
// workspace. An empty image clears the override (NULL = the runner's default
// windshift-agent image) (WI-450).
Expand Down
18 changes: 15 additions & 3 deletions internal/services/agent_profile_service.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,10 @@ type ProfileValidationError struct {
type ProfileValidationResult struct {
Ready bool `json:"ready"`
Errors []ProfileValidationError `json:"errors"`
// ProfileVersion is the definition version the checks ran against.
// Activation only succeeds while the stored version still matches it, so a
// configuration change during validation cannot be marked Ready.
ProfileVersion int `json:"profile_version"`
}

// UpdateStudioProfileRequest carries the mutable overview fields. Profile type
Expand Down Expand Up @@ -526,7 +530,7 @@ func (s *BindingService) ValidateStudioProfile(ctx context.Context, workspaceID,
return nil, fmt.Errorf("load agent profile: %w", err)
}

result := &ProfileValidationResult{Errors: []ProfileValidationError{}}
result := &ProfileValidationResult{Errors: []ProfileValidationError{}, ProfileVersion: binding.ProfileVersion}
var active, isAgent bool
if err := s.db.QueryRowContext(ctx, `
SELECT COALESCE(is_active, false), COALESCE(is_agent, false)
Expand Down Expand Up @@ -649,15 +653,23 @@ func (s *BindingService) ActivateStudioProfile(ctx context.Context, workspaceID,
if err != nil {
return nil, nil, err
}
return s.activateStudioProfile(ctx, workspaceID, bindingID, validation)
}

// activateStudioProfile transitions a validated definition to Ready, but only
// while the stored profile version still matches the one the validation ran
// against. A concurrent configuration change bumps the version, so this refuses
// to mark a changed definition Ready using validation of an older one.
func (s *BindingService) activateStudioProfile(ctx context.Context, workspaceID, bindingID int, validation *ProfileValidationResult) (*models.WorkspaceAgentBinding, *ProfileValidationResult, error) {
if !validation.Ready {
return nil, validation, ErrAgentProfileValidationFailed
}
rows, err := s.repo.SetLifecycle(ctx, bindingID, workspaceID, models.AgentLifecycleReady)
rows, err := s.repo.SetLifecycleIfVersion(ctx, bindingID, workspaceID, models.AgentLifecycleReady, validation.ProfileVersion)
if err != nil {
return nil, validation, err
}
if rows == 0 {
return nil, validation, ErrBindingUnavailable
return nil, validation, ErrAgentProfileVersionConflict
}
binding, err := s.repo.Get(ctx, bindingID)
if err != nil {
Expand Down
19 changes: 8 additions & 11 deletions internal/services/binding_service.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -723,18 +723,15 @@ func (s *BindingService) UpdateAgentConfig(ctx context.Context, workspaceID, bin
return ErrLLMConnectionInvalid
}
}
if err := s.repo.UpdateInstructions(ctx, bindingID, workspaceID, instructions); err != nil {
return err
}
if llmConnectionID != nil {
if err := s.repo.UpdateLLMConnection(ctx, bindingID, workspaceID, *llmConnectionID); err != nil {
return err
}
}
var runnerImageArg *string
if runnerImage != nil {
if err := s.repo.UpdateRunnerImage(ctx, bindingID, workspaceID, newRunnerImage); err != nil {
return err
}
runnerImageArg = &newRunnerImage
}
// One atomic write so the model connection, readiness invalidation, and the
// single version bump land together; a concurrent activation must never see
// the old connection with a fresh version.
if err := s.repo.UpdateAgentConfig(ctx, bindingID, workspaceID, instructions, llmConnectionID, runnerImageArg); err != nil {
return err
}
if s.skills == nil {
if len(skillIDs) > 0 {
Expand Down
Loading