change: use cleanup handler for runstate instead of finalizer on a run (

#661)
obot-platform · Nov 25, 2024 · 1ed16ce · 1ed16ce
1 parent 6254fe3
commit 1ed16ce
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 16 deletions.
diff --git a/pkg/controller/handlers/runs/runs.go b/pkg/controller/handlers/runs/runs.go
@@ -2,15 +2,14 @@ package runs
 
 import (
 	"fmt"
+	"slices"
 
 	"github.com/gptscript-ai/go-gptscript"
 	"github.com/otto8-ai/nah/pkg/router"
 	"github.com/otto8-ai/otto8/logger"
 	"github.com/otto8-ai/otto8/pkg/invoke"
 	v1 "github.com/otto8-ai/otto8/pkg/storage/apis/otto.otto8.ai/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
 var log = logger.Package()
@@ -23,16 +22,6 @@ func New(invoker *invoke.Invoker) *Handler {
 	return &Handler{invoker: invoker}
 }
 
-func (*Handler) DeleteRunState(req router.Request, resp router.Response) error {
-	run := req.Object.(*v1.Run)
-	return client.IgnoreNotFound(req.Delete(&v1.RunState{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      run.Name,
-			Namespace: run.Namespace,
-		},
-	}))
-}
-
 func (h *Handler) Resume(req router.Request, _ router.Response) error {
 	run := req.Object.(*v1.Run)
 	var thread v1.Thread
@@ -70,3 +59,21 @@ func (h *Handler) Resume(req router.Request, _ router.Response) error {
 
 	return h.invoker.Resume(req.Ctx, req.Client, &thread, run)
 }
+
+// MigrateRemoveRunFinalizer (to be removed) removes the run finalizer from the run object which was used to cascade delete the run state,
+// which was moved to its own cleanup handler.
+func (h *Handler) MigrateRemoveRunFinalizer(req router.Request, _ router.Response) error {
+	run := req.Object.(*v1.Run)
+	changed := false
+	run.Finalizers = slices.DeleteFunc(run.ObjectMeta.Finalizers, func(i string) bool {
+		if i == v1.RunFinalizer {
+			changed = true
+			return true
+		}
+		return false
+	})
+	if changed {
+		return req.Client.Update(req.Ctx, run)
+	}
+	return nil
+}
diff --git a/pkg/controller/routes.go b/pkg/controller/routes.go
@@ -37,10 +37,13 @@ func (c *Controller) setupRoutes() error {
 	oauthLogins := oauthapp.NewLogin(c.services.Invoker, c.services.ServerURL)
 
 	// Runs
-	root.Type(&v1.Run{}).FinalizeFunc(v1.RunFinalizer, runs.DeleteRunState)
+	root.Type(&v1.Run{}).HandlerFunc(runs.MigrateRemoveRunFinalizer) // to be removed
 	root.Type(&v1.Run{}).HandlerFunc(cleanup.Cleanup)
 	root.Type(&v1.Run{}).HandlerFunc(runs.Resume)
 
+	// RunStates
+	root.Type(&v1.RunState{}).HandlerFunc(cleanup.Cleanup)
+
 	// Threads
 	root.Type(&v1.Thread{}).HandlerFunc(cleanup.Cleanup)
 	root.Type(&v1.Thread{}).HandlerFunc(threads.CreateWorkspaces)

diff --git a/pkg/invoke/invoker.go b/pkg/invoke/invoker.go
@@ -345,7 +345,6 @@ func (i *Invoker) createRun(ctx context.Context, c kclient.WithWatch, thread *v1
 		ObjectMeta: metav1.ObjectMeta{
 			GenerateName: system.RunPrefix,
 			Namespace:    thread.Namespace,
-			Finalizers:   []string{v1.RunFinalizer},
 		},
 		Spec: v1.RunSpec{
 			Synchronous:           opts.Synchronous,

diff --git a/pkg/storage/apis/otto.otto8.ai/v1/run.go b/pkg/storage/apis/otto.otto8.ai/v1/run.go
@@ -7,7 +7,7 @@ import (
 )
 
 const (
-	RunFinalizer               = "otto.otto8.ai/run"
+	RunFinalizer               = "otto.otto8.ai/run" // to be removed
 	KnowledgeFileFinalizer     = "otto.otto8.ai/knowledge-file"
 	WorkspaceFinalizer         = "otto.otto8.ai/workspace"
 	KnowledgeSetFinalizer      = "otto.otto8.ai/knowledge-set"

diff --git a/pkg/storage/apis/otto.otto8.ai/v1/runstate.go b/pkg/storage/apis/otto.otto8.ai/v1/runstate.go
@@ -25,7 +25,9 @@ type RunStateSpec struct {
 }
 
 func (in *RunState) DeleteRefs() []Ref {
-	return []Ref{}
+	return []Ref{
+		{ObjType: &Run{}, Name: in.Name, Namespace: in.Namespace},
+	}
 }
 
 // +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object