From 527ba212b48fc164a9ac8c780ae19687be25d0c6 Mon Sep 17 00:00:00 2001 From: Donnie Adams Date: Wed, 29 Jan 2025 22:05:48 -0500 Subject: [PATCH] fix: remove agent and workflow credentials on delete Signed-off-by: Donnie Adams --- pkg/api/handlers/agent.go | 1 + pkg/api/handlers/workflows.go | 1 + pkg/controller/data/agent.yaml | 2 ++ .../handlers/cleanup/credentials.go | 33 +++++++++++++++++++ pkg/controller/handlers/threads/threads.go | 19 ----------- pkg/controller/routes.go | 5 ++- pkg/storage/apis/obot.obot.ai/v1/run.go | 3 ++ pkg/storage/apis/obot.obot.ai/v1/thread.go | 4 --- 8 files changed, 44 insertions(+), 24 deletions(-) create mode 100644 pkg/controller/handlers/cleanup/credentials.go diff --git a/pkg/api/handlers/agent.go b/pkg/api/handlers/agent.go index 49b7f350a..d18b428e0 100644 --- a/pkg/api/handlers/agent.go +++ b/pkg/api/handlers/agent.go @@ -181,6 +181,7 @@ func (a *AgentHandler) Create(req api.Context) error { ObjectMeta: metav1.ObjectMeta{ GenerateName: system.AgentPrefix, Namespace: req.Namespace(), + Finalizers: []string{v1.AgentFinalizer}, }, Spec: v1.AgentSpec{ Manifest: manifest, diff --git a/pkg/api/handlers/workflows.go b/pkg/api/handlers/workflows.go index 54a471323..aa993caaf 100644 --- a/pkg/api/handlers/workflows.go +++ b/pkg/api/handlers/workflows.go @@ -234,6 +234,7 @@ func (a *WorkflowHandler) Create(req api.Context) error { ObjectMeta: metav1.ObjectMeta{ GenerateName: system.WorkflowPrefix, Namespace: req.Namespace(), + Finalizers: []string{v1.WorkflowFinalizer}, }, Spec: v1.WorkflowSpec{ Manifest: manifest, diff --git a/pkg/controller/data/agent.yaml b/pkg/controller/data/agent.yaml index c2fe9fcbe..cd64210b5 100644 --- a/pkg/controller/data/agent.yaml +++ b/pkg/controller/data/agent.yaml @@ -3,6 +3,8 @@ kind: Agent metadata: name: a1-obot namespace: default + finalizers: + - obot.obot.ai/agent spec: manifest: name: Obot diff --git a/pkg/controller/handlers/cleanup/credentials.go b/pkg/controller/handlers/cleanup/credentials.go new file mode 100644 index 000000000..ae55b4b9b --- /dev/null +++ b/pkg/controller/handlers/cleanup/credentials.go @@ -0,0 +1,33 @@ +package cleanup + +import ( + "github.com/gptscript-ai/go-gptscript" + "github.com/obot-platform/nah/pkg/router" +) + +type Credentials struct { + gClient *gptscript.GPTScript +} + +func NewCredentials(gClient *gptscript.GPTScript) *Credentials { + return &Credentials{ + gClient: gClient, + } +} + +func (c *Credentials) Remove(req router.Request, _ router.Response) error { + creds, err := c.gClient.ListCredentials(req.Ctx, gptscript.ListCredentialsOptions{ + CredentialContexts: []string{req.Object.GetName()}, + }) + if err != nil { + return err + } + + for _, cred := range creds { + if err := c.gClient.DeleteCredential(req.Ctx, req.Object.GetName(), cred.ToolName); err != nil { + return err + } + } + + return nil +} diff --git a/pkg/controller/handlers/threads/threads.go b/pkg/controller/handlers/threads/threads.go index 985d15c63..a0bc1da2a 100644 --- a/pkg/controller/handlers/threads/threads.go +++ b/pkg/controller/handlers/threads/threads.go @@ -93,25 +93,6 @@ func (t *Handler) CreateWorkspaces(req router.Request, _ router.Response) error return nil } -func (t *Handler) CleanupThread(req router.Request, _ router.Response) error { - thread := req.Object.(*v1.Thread) - - creds, err := t.gptScript.ListCredentials(req.Ctx, gptscript.ListCredentialsOptions{ - CredentialContexts: []string{thread.Name}, - }) - if err != nil { - return err - } - - for _, cred := range creds { - if err := t.gptScript.DeleteCredential(req.Ctx, thread.Name, cred.ToolName); err != nil { - return err - } - } - - return nil -} - func (t *Handler) CreateKnowledgeSet(req router.Request, _ router.Response) error { thread := req.Object.(*v1.Thread) if len(thread.Status.KnowledgeSetNames) > 0 || thread.Spec.AgentName == "" { diff --git a/pkg/controller/routes.go b/pkg/controller/routes.go index acc38b304..b981ff3d7 100644 --- a/pkg/controller/routes.go +++ b/pkg/controller/routes.go @@ -46,6 +46,7 @@ func (c *Controller) setupRoutes() error { knowledgesummary := knowledgesummary.NewHandler(c.services.GPTClient) toolInfo := toolinfo.New(c.services.GPTClient) threads := threads.NewHandler(c.services.GPTClient) + credentialCleanup := cleanup.NewCredentials(c.services.GPTClient) // Runs root.Type(&v1.Run{}).HandlerFunc(removeOldFinalizers) @@ -61,7 +62,7 @@ func (c *Controller) setupRoutes() error { root.Type(&v1.Thread{}).HandlerFunc(threads.CreateKnowledgeSet) root.Type(&v1.Thread{}).HandlerFunc(threads.WorkflowState) root.Type(&v1.Thread{}).HandlerFunc(knowledgesummary.Summarize) - root.Type(&v1.Thread{}).FinalizeFunc(v1.ThreadFinalizer, threads.CleanupThread) + root.Type(&v1.Thread{}).FinalizeFunc(v1.ThreadFinalizer, credentialCleanup.Remove) // KnowledgeSummary root.Type(&v1.KnowledgeSummary{}).HandlerFunc(cleanup.Cleanup) @@ -74,6 +75,7 @@ func (c *Controller) setupRoutes() error { root.Type(&v1.Workflow{}).HandlerFunc(alias.AssignAlias) root.Type(&v1.Workflow{}).HandlerFunc(toolInfo.SetToolInfoStatus) root.Type(&v1.Workflow{}).HandlerFunc(generationed.UpdateObservedGeneration) + root.Type(&v1.Workflow{}).FinalizeFunc(v1.WorkflowFinalizer, credentialCleanup.Remove) // WorkflowExecutions root.Type(&v1.WorkflowExecution{}).HandlerFunc(cleanup.Cleanup) @@ -86,6 +88,7 @@ func (c *Controller) setupRoutes() error { root.Type(&v1.Agent{}).HandlerFunc(alias.AssignAlias) root.Type(&v1.Agent{}).HandlerFunc(toolInfo.SetToolInfoStatus) root.Type(&v1.Agent{}).HandlerFunc(generationed.UpdateObservedGeneration) + root.Type(&v1.Agent{}).FinalizeFunc(v1.AgentFinalizer, credentialCleanup.Remove) // Uploads root.Type(&v1.KnowledgeSource{}).HandlerFunc(cleanup.Cleanup) diff --git a/pkg/storage/apis/obot.obot.ai/v1/run.go b/pkg/storage/apis/obot.obot.ai/v1/run.go index c3bd94e22..f832e72a5 100644 --- a/pkg/storage/apis/obot.obot.ai/v1/run.go +++ b/pkg/storage/apis/obot.obot.ai/v1/run.go @@ -8,11 +8,14 @@ import ( const ( RunFinalizer = "obot.obot.ai/run" + ThreadFinalizer = "obot.obot.ai/thread" KnowledgeFileFinalizer = "obot.obot.ai/knowledge-file" WorkspaceFinalizer = "obot.obot.ai/workspace" KnowledgeSetFinalizer = "obot.obot.ai/knowledge-set" KnowledgeSourceFinalizer = "obot.obot.ai/knowledge-source" ToolReferenceFinalizer = "obot.obot.ai/tool-reference" + AgentFinalizer = "obot.obot.ai/agent" + WorkflowFinalizer = "obot.obot.ai/workflow" ModelProviderSyncAnnotation = "obot.ai/model-provider-sync" WorkflowSyncAnnotation = "obot.ai/workflow-sync" diff --git a/pkg/storage/apis/obot.obot.ai/v1/thread.go b/pkg/storage/apis/obot.obot.ai/v1/thread.go index 427d87ace..a3470b002 100644 --- a/pkg/storage/apis/obot.obot.ai/v1/thread.go +++ b/pkg/storage/apis/obot.obot.ai/v1/thread.go @@ -8,10 +8,6 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) -const ( - ThreadFinalizer = "obot.obot.ai/thread" -) - // +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object type Thread struct {