fix: stop looping if error occurs for agent login

Signed-off-by: Donnie Adams <[email protected]>
obot-platform · Nov 5, 2024 · 4264e5a · 4264e5a
1 parent bd1ba75
commit 4264e5a
Show file tree

Hide file tree

Showing 7 changed files with 73 additions and 47 deletions.
diff --git a/apiclient/types/zz_generated.deepcopy.go b/apiclient/types/zz_generated.deepcopy.go
diff --git a/pkg/api/handlers/agent.go b/pkg/api/handlers/agent.go
@@ -487,7 +487,7 @@ func (a *AgentHandler) EnsureCredentialForKnowledgeSource(req api.Context) error
 	}
 
 	oauthLogin, err = wait.For(req.Context(), req.Storage, oauthLogin, func(obj *v1.OAuthAppLogin) bool {
-		return obj.Status.Authenticated || obj.Status.Error != "" || obj.Status.URL != ""
+		return obj.Status.External.Authenticated || obj.Status.External.Error != "" || obj.Status.External.URL != ""
 	}, wait.Option{
 		Create: true,
 	})
@@ -499,7 +499,7 @@ func (a *AgentHandler) EnsureCredentialForKnowledgeSource(req api.Context) error
 	if agent.Status.External.AuthStatus == nil {
 		agent.Status.External.AuthStatus = make(map[string]types.OAuthAppLoginAuthStatus)
 	}
-	agent.Status.External.AuthStatus[ref] = oauthLogin.Status.OAuthAppLoginAuthStatus
+	agent.Status.External.AuthStatus[ref] = oauthLogin.Status.External
 	return req.Write(convertAgent(agent, server.GetURLPrefix(req)))
 }
 

diff --git a/pkg/controller/handlers/agents/agents.go b/pkg/controller/handlers/agents/agents.go
@@ -13,6 +13,7 @@ import (
 	"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"
 )
 
@@ -77,41 +78,45 @@ func BackPopulateAuthStatus(req router.Request, _ router.Response) error {
 	agent := req.Object.(*v1.Agent)
 
 	var logins v1.OAuthAppLoginList
-	if err := req.List(&logins, &kclient.ListOptions{Namespace: agent.Namespace}); err != nil {
+	if err := req.List(&logins, &kclient.ListOptions{
+		FieldSelector: fields.SelectorFromSet(map[string]string{"spec.credentialContext": agent.Name}),
+		Namespace:     agent.Namespace,
+	}); err != nil {
 		return err
 	}
 
 	for _, login := range logins.Items {
-		if login.Status.Authenticated || (login.Status.Required != nil && !*login.Status.Required) || login.Spec.ToolReference == "" {
+		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 {
-			login.Status.Error = fmt.Sprintf("failed to get credential tool for knowledge source [%s]: %v", agent.Name, err)
+			login.Status.External.Error = fmt.Sprintf("failed to get credential tool for knowledge source [%s]: %v", agent.Name, err)
+		} else {
+			required = &[]bool{credentialTool != ""}[0]
+			updateRequired = updateRequired || login.Status.External.Required == nil || *login.Status.External.Required != *required
+			login.Status.External.Required = required
 		}
 
-		required := credentialTool != ""
-		updateRequired = updateRequired || login.Status.Required == nil || *login.Status.Required != required
-		login.Status.Required = &required
-
-		if required {
+		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.Error != ""
-				login.Status.Error = ""
+				updateRequired = updateRequired || login.Status.External.Error != ""
+				login.Status.External.Error = ""
 			} else if err != nil {
-				login.Status.Error = fmt.Sprintf("failed to get oauth app login for agent [%s]: %v", agent.Name, err)
+				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, oauthAppLogin.Status.OAuthAppLoginAuthStatus)
-				login.Status.OAuthAppLoginAuthStatus = oauthAppLogin.Status.OAuthAppLoginAuthStatus
+				updateRequired = updateRequired || equality.Semantic.DeepEqual(login.Status.External, oauthAppLogin.Status.External)
+				login.Status.External = oauthAppLogin.Status.External
 			}
 		}
 
 		if agent.Status.External.AuthStatus == nil {
 			agent.Status.External.AuthStatus = make(map[string]types.OAuthAppLoginAuthStatus)
 		}
-		agent.Status.External.AuthStatus[login.Spec.ToolReference] = login.Status.OAuthAppLoginAuthStatus
+		agent.Status.External.AuthStatus[login.Spec.ToolReference] = login.Status.External
 	}
 
 	if updateRequired {

diff --git a/pkg/controller/handlers/oauthapp/oauthapplogin.go b/pkg/controller/handlers/oauthapp/oauthapplogin.go
@@ -27,7 +27,7 @@ func NewLogin(invoker *invoke.Invoker) *LoginHandler {
 
 func (h *LoginHandler) RunTool(req router.Request, _ router.Response) error {
 	login := req.Object.(*v1.OAuthAppLogin)
-	if login.Status.Authenticated || login.Status.Error != "" || login.Spec.ToolReference == "" {
+	if login.Status.External.Authenticated || login.Status.External.Error != "" || login.Spec.ToolReference == "" {
 		return nil
 	}
 
@@ -70,13 +70,13 @@ func (h *LoginHandler) RunTool(req router.Request, _ router.Response) error {
 	for frame := range task.Events {
 		if frame.Prompt != nil && frame.Prompt.Metadata["authURL"] != "" {
 			if err = updateLoginExternalStatus(req.Ctx, req.Client, login, v1.OAuthAppLoginStatus{
-				OAuthAppLoginAuthStatus: types.OAuthAppLoginAuthStatus{
+				External: types.OAuthAppLoginAuthStatus{
 					URL:      frame.Prompt.Metadata["authURL"],
 					Required: &[]bool{true}[0],
 				},
 			}); err != nil {
 				if setErrorErr := updateLoginExternalStatus(req.Ctx, req.Client, login, v1.OAuthAppLoginStatus{
-					OAuthAppLoginAuthStatus: types.OAuthAppLoginAuthStatus{
+					External: types.OAuthAppLoginAuthStatus{
 						Error: err.Error(),
 					},
 				}); setErrorErr != nil {
@@ -94,7 +94,7 @@ func (h *LoginHandler) RunTool(req router.Request, _ router.Response) error {
 	}
 
 	return updateLoginExternalStatus(req.Ctx, req.Client, login, v1.OAuthAppLoginStatus{
-		OAuthAppLoginAuthStatus: types.OAuthAppLoginAuthStatus{
+		External: types.OAuthAppLoginAuthStatus{
 			Error:         errMessage,
 			Authenticated: errMessage == "",
 			URL:           "",

diff --git a/pkg/storage/apis/otto.otto8.ai/v1/oauthapp.go b/pkg/storage/apis/otto.otto8.ai/v1/oauthapp.go
@@ -163,7 +163,7 @@ type OAuthAppLoginSpec struct {
 }
 
 type OAuthAppLoginStatus struct {
-	types.OAuthAppLoginAuthStatus `json:",inline"`
+	External types.OAuthAppLoginAuthStatus `json:"external,omitempty"`
 }
 
 // +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object

diff --git a/pkg/storage/apis/otto.otto8.ai/v1/zz_generated.deepcopy.go b/pkg/storage/apis/otto.otto8.ai/v1/zz_generated.deepcopy.go
diff --git a/pkg/storage/openapi/generated/openapi_generated.go b/pkg/storage/openapi/generated/openapi_generated.go