fix: move workflow migration back to the controller (#1795)

Signed-off-by: Grant Linville <[email protected]>

Author: g-linville

pkg/controller/controller.go

@@ -7,7 +7,6 @@ import (
 	"github.com/obot-platform/nah/pkg/router"
 	"github.com/obot-platform/obot/pkg/controller/data"
 	"github.com/obot-platform/obot/pkg/controller/handlers/toolreference"
-	"github.com/obot-platform/obot/pkg/controller/handlers/workflow"
 	"github.com/obot-platform/obot/pkg/services"
 	kclient "sigs.k8s.io/controller-runtime/pkg/client"
 
@@ -44,9 +43,6 @@ func (c *Controller) PreStart(ctx context.Context) error {
 	if err := toolreference.MigrateToolNames(ctx, c.services.StorageClient); err != nil {
 		return fmt.Errorf("failed to migrate tool names: %w", err)
 	}
-	if err := workflow.SetAdditionalCredentialContexts(ctx, c.services.StorageClient); err != nil {
-		return fmt.Errorf("failed to set workflow additional credential contexts: %w", err)
-	}
 	return nil
 }
 

pkg/controller/handlers/workflow/migration.go

@@ -6,6 +6,7 @@ import (
 	"github.com/obot-platform/nah/pkg/apply"
 	"github.com/obot-platform/nah/pkg/router"
 	v1 "github.com/obot-platform/obot/pkg/storage/apis/obot.obot.ai/v1"
+	kclient "sigs.k8s.io/controller-runtime/pkg/client"
 )
 
 func removeOldFinalizers(req router.Request, _ router.Response) error {
@@ -56,3 +57,27 @@ func changeWorkflowStepOwnerGVK(req router.Request, _ router.Response) error {
 	}
 	return nil
 }
+
+func setWorkflowAdditionalCredentialContexts(req router.Request, _ router.Response) error {
+	wf := req.Object.(*v1.Workflow)
+
+	if len(wf.Spec.AdditionalCredentialContexts) != 0 || wf.Spec.ThreadName == "" {
+		return nil
+	}
+
+	var thread v1.Thread
+	if err := req.Client.Get(req.Ctx, kclient.ObjectKey{Namespace: wf.Namespace, Name: wf.Spec.ThreadName}, &thread); err != nil {
+		return err
+	}
+
+	if thread.Spec.AgentName == "" {
+		return nil
+	}
+
+	wf.Spec.AdditionalCredentialContexts = []string{thread.Spec.AgentName}
+	if err := req.Client.Update(req.Ctx, wf); err != nil {
+		return err
+	}
+
+	return nil
+}

pkg/controller/migrations.go

@@ -79,6 +79,7 @@ func (c *Controller) setupRoutes() error {
 	root.Type(&v1.Workflow{}).HandlerFunc(toolInfo.SetToolInfoStatus)
 	root.Type(&v1.Workflow{}).HandlerFunc(toolInfo.RemoveUnneededCredentials)
 	root.Type(&v1.Workflow{}).HandlerFunc(generationed.UpdateObservedGeneration)
+	root.Type(&v1.Workflow{}).HandlerFunc(setWorkflowAdditionalCredentialContexts)
 	root.Type(&v1.Workflow{}).FinalizeFunc(v1.WorkflowFinalizer, credentialCleanup.Remove)
 
 	// WorkflowExecutions