fix: more reliably set the oauth login status and remove retry logic

Signed-off-by: Donnie Adams <[email protected]>
otto8-ai · Nov 22, 2024 · f5f9192 · f5f9192
1 parent b920c70
commit f5f9192
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 40 deletions.
diff --git a/pkg/controller/handlers/agents/agents.go b/pkg/controller/handlers/agents/agents.go
@@ -11,7 +11,6 @@ import (
 	v1 "github.com/otto8-ai/otto8/pkg/storage/apis/otto.otto8.ai/v1"
 	"github.com/otto8-ai/otto8/pkg/system"
 	"k8s.io/apimachinery/pkg/api/equality"
-	apierror "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/fields"
 	kclient "sigs.k8s.io/controller-runtime/pkg/client"
@@ -74,7 +73,6 @@ func CreateWorkspaceAndKnowledgeSet(req router.Request, _ router.Response) error
 }
 
 func BackPopulateAuthStatus(req router.Request, _ router.Response) error {
-	var updateRequired bool
 	agent := req.Object.(*v1.Agent)
 
 	var logins v1.OAuthAppLoginList
@@ -85,11 +83,10 @@ func BackPopulateAuthStatus(req router.Request, _ router.Response) error {
 		return err
 	}
 
+	existingLoginInfo := agent.Status.AuthStatus
+	updateRequired := len(existingLoginInfo) != len(logins.Items)
+	agent.Status.AuthStatus = make(map[string]types.OAuthAppLoginAuthStatus)
 	for _, login := range logins.Items {
-		if login.Status.External.Authenticated || (login.Status.External.Required != nil && !*login.Status.External.Required) || login.Spec.ToolReference == "" {
-			continue
-		}
-
 		var required *bool
 		credentialTool, err := v1.CredentialTool(req.Ctx, req.Client, agent.Namespace, login.Spec.ToolReference)
 		if err != nil {
@@ -101,21 +98,9 @@ func BackPopulateAuthStatus(req router.Request, _ router.Response) error {
 		}
 
 		if required != nil && *required {
-			var oauthAppLogin v1.OAuthAppLogin
-			if err = req.Get(&oauthAppLogin, agent.Namespace, system.OAuthAppLoginPrefix+agent.Name+login.Spec.ToolReference); apierror.IsNotFound(err) {
-				updateRequired = updateRequired || login.Status.External.Error != ""
-				login.Status.External.Error = ""
-			} else if err != nil {
-				login.Status.External.Error = fmt.Sprintf("failed to get oauth app login for agent [%s]: %v", agent.Name, err)
-			} else {
-				updateRequired = updateRequired || equality.Semantic.DeepEqual(login.Status.External, oauthAppLogin.Status.External)
-				login.Status.External = oauthAppLogin.Status.External
-			}
+			updateRequired = updateRequired || equality.Semantic.DeepEqual(login.Status.External, existingLoginInfo[login.Spec.ToolReference])
 		}
 
-		if agent.Status.AuthStatus == nil {
-			agent.Status.AuthStatus = make(map[string]types.OAuthAppLoginAuthStatus)
-		}
 		agent.Status.AuthStatus[login.Spec.ToolReference] = login.Status.External
 	}
 

diff --git a/pkg/controller/handlers/oauthapp/oauthapplogin.go b/pkg/controller/handlers/oauthapp/oauthapplogin.go
@@ -1,7 +1,6 @@
 package oauthapp
 
 import (
-	"context"
 	"errors"
 	"time"
 
@@ -14,8 +13,6 @@ import (
 	"github.com/otto8-ai/otto8/pkg/system"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/client-go/util/retry"
-	kclient "sigs.k8s.io/controller-runtime/pkg/client"
 )
 
 type LoginHandler struct {
@@ -75,12 +72,13 @@ func (h *LoginHandler) RunTool(req router.Request, _ router.Response) error {
 	// Ensure the task is stopped when this handler returns.
 	defer task.Close()
 
-	if err = updateLoginExternalStatus(req.Ctx, req.Client, login, v1.OAuthAppLoginStatus{}); err != nil {
+	login.Status.External = types.OAuthAppLoginAuthStatus{}
+	if err = req.Client.Status().Update(req.Ctx, login); err != nil {
 		return err
 	}
 
 	originalUID := login.UID
-	tick := time.NewTicker(5 * time.Second)
+	tick := time.NewTicker(time.Second)
 	defer tick.Stop()
 
 outer:
@@ -99,17 +97,19 @@ outer:
 			}
 
 			if frame.Prompt != nil && frame.Prompt.Metadata["authURL"] != "" {
-				if err = updateLoginExternalStatus(req.Ctx, req.Client, login, v1.OAuthAppLoginStatus{
+				login.Status = v1.OAuthAppLoginStatus{
 					External: types.OAuthAppLoginAuthStatus{
 						URL:      frame.Prompt.Metadata["authURL"],
 						Required: &[]bool{true}[0],
 					},
-				}); err != nil {
-					if setErrorErr := updateLoginExternalStatus(req.Ctx, req.Client, login, v1.OAuthAppLoginStatus{
+				}
+				if err = req.Client.Status().Update(req.Ctx, login); err != nil {
+					login.Status = v1.OAuthAppLoginStatus{
 						External: types.OAuthAppLoginAuthStatus{
 							Error: err.Error(),
 						},
-					}); setErrorErr != nil {
+					}
+					if setErrorErr := req.Client.Status().Update(req.Ctx, login); setErrorErr != nil {
 						err = errors.Join(err, setErrorErr)
 					}
 					return err
@@ -126,23 +126,14 @@ outer:
 		errMessage = err.Error()
 	}
 
-	return updateLoginExternalStatus(req.Ctx, req.Client, login, v1.OAuthAppLoginStatus{
+	login.Status = v1.OAuthAppLoginStatus{
 		External: types.OAuthAppLoginAuthStatus{
 			Error:         errMessage,
 			Authenticated: errMessage == "",
 			URL:           "",
 			Required:      &[]bool{true}[0],
 		},
-	})
-}
-
-func updateLoginExternalStatus(ctx context.Context, client kclient.Client, login *v1.OAuthAppLogin, status v1.OAuthAppLoginStatus) error {
-	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
-		if err := client.Get(ctx, router.Key(login.Namespace, login.Name), login); err != nil {
-			return err
-		}
+	}
 
-		login.Status = status
-		return client.Status().Update(ctx, login)
-	})
+	return req.Client.Status().Update(req.Ctx, login)
 }