From 9a4917b7ad1c0c6224b4a80a97aa7b6c370c924c Mon Sep 17 00:00:00 2001
From: Thorsten Klein <tk@thklein.io>
Date: Thu, 5 Dec 2024 14:47:12 +0100
Subject: [PATCH] revert 1ed16ce... change: use cleanup handler for runstate
 instead of finalizer on a run (#661)

---
 pkg/controller/handlers/runs/runs.go          | 31 +++++++------------
 pkg/controller/routes.go                      |  5 +--
 pkg/invoke/invoker.go                         |  1 +
 pkg/storage/apis/otto.otto8.ai/v1/run.go      |  2 +-
 pkg/storage/apis/otto.otto8.ai/v1/runstate.go |  4 +--
 5 files changed, 16 insertions(+), 27 deletions(-)

diff --git a/pkg/controller/handlers/runs/runs.go b/pkg/controller/handlers/runs/runs.go
index d9c02f94..61086f33 100644
--- a/pkg/controller/handlers/runs/runs.go
+++ b/pkg/controller/handlers/runs/runs.go
@@ -2,7 +2,6 @@ package runs
 
 import (
 	"fmt"
-	"slices"
 	"time"
 
 	"github.com/gptscript-ai/go-gptscript"
@@ -11,6 +10,8 @@ import (
 	"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,6 +24,16 @@ 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
@@ -75,21 +86,3 @@ func (h *Handler) DeleteFinished(req router.Request, _ router.Response) error {
 	}
 	return nil
 }
-
-// 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
index 66f1ab9e..2cf2eaf5 100644
--- a/pkg/controller/routes.go
+++ b/pkg/controller/routes.go
@@ -37,14 +37,11 @@ func (c *Controller) setupRoutes() error {
 	oauthLogins := oauthapp.NewLogin(c.services.Invoker, c.services.ServerURL)
 
 	// Runs
-	root.Type(&v1.Run{}).HandlerFunc(runs.MigrateRemoveRunFinalizer) // to be removed
+	root.Type(&v1.Run{}).FinalizeFunc(v1.RunFinalizer, runs.DeleteRunState)
 	root.Type(&v1.Run{}).HandlerFunc(runs.DeleteFinished)
 	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
index 1dd1dbed..81e3431c 100644
--- a/pkg/invoke/invoker.go
+++ b/pkg/invoke/invoker.go
@@ -359,6 +359,7 @@ 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
index 077229ef..013c2c82 100644
--- 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" // to be removed
+	RunFinalizer             = "otto.otto8.ai/run"
 	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
index d39709c4..292f239e 100644
--- a/pkg/storage/apis/otto.otto8.ai/v1/runstate.go
+++ b/pkg/storage/apis/otto.otto8.ai/v1/runstate.go
@@ -25,9 +25,7 @@ type RunStateSpec struct {
 }
 
 func (in *RunState) DeleteRefs() []Ref {
-	return []Ref{
-		{ObjType: &Run{}, Name: in.Name, Namespace: in.Namespace},
-	}
+	return []Ref{}
 }
 
 // +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object