feat: support PATs for authenticating with knowledge sources

Signed-off-by: Donnie Adams <[email protected]>
obot-platform · Jan 22, 2025 · 1b90511 · 1b90511
1 parent fc9770e
commit 1b90511
Show file tree

Hide file tree

Showing 8 changed files with 135 additions and 31 deletions.
diff --git a/apiclient/types/oauthapp.go b/apiclient/types/oauthapp.go
@@ -45,8 +45,11 @@ type OAuthAppManifest struct {
 type OAuthAppList List[OAuthApp]
 
 type OAuthAppLoginAuthStatus struct {
-	URL           string `json:"url,omitempty"`
-	Authenticated bool   `json:"authenticated,omitempty"`
-	Required      *bool  `json:"required,omitempty"`
-	Error         string `json:"error,omitempty"`
+	URL             string   `json:"url,omitempty"`
+	PromptMessage   string   `json:"promptMessage,omitempty"`
+	PromptFields    []string `json:"promptFields,omitempty"`
+	PromptSensitive bool     `json:"promptSensitive,omitempty"`
+	Authenticated   bool     `json:"authenticated,omitempty"`
+	Required        *bool    `json:"required,omitempty"`
+	Error           string   `json:"error,omitempty"`
 }
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
@@ -734,12 +734,16 @@ func (a *AgentHandler) EnsureCredentialForKnowledgeSource(req api.Context) error
 		return req.WriteCreated(resp)
 	}
 
-	credentialTools, err := v1.CredentialTools(req.Context(), req.Storage, req.Namespace(), ref)
-	if err != nil {
+	var toolReference v1.ToolReference
+	if err := req.Get(&toolReference, ref); err != nil {
 		return err
 	}
 
-	if len(credentialTools) == 0 {
+	if toolReference.Status.Tool == nil {
+		return types.NewErrHttp(http.StatusTooEarly, fmt.Sprintf("tool %q is not processed yet", ref))
+	}
+
+	if len(toolReference.Status.Tool.Credentials) == 0 {
 		// The only way to get here is if the controller hasn't set the field yet.
 		if agent.Status.AuthStatus == nil {
 			agent.Status.AuthStatus = make(map[string]types.OAuthAppLoginAuthStatus)
@@ -754,24 +758,39 @@ func (a *AgentHandler) EnsureCredentialForKnowledgeSource(req api.Context) error
 		return req.WriteCreated(resp)
 	}
 
+	var supportedPATOAuthApps []string
+	if toolReference.Status.Tool.Metadata["oauthPATSupported"] == "true" {
+		supportedPATOAuthApps = append(supportedPATOAuthApps, toolReference.Status.Tool.Metadata["oauth"])
+	}
+
+	var oauthApps []string
+	if req.URL.Query().Get("pat") != "true" {
+		oauthApps = append(oauthApps, toolReference.Status.Tool.Metadata["oauth"])
+	}
+
 	oauthLogin := &v1.OAuthAppLogin{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      system.OAuthAppLoginPrefix + agent.Name + ref,
 			Namespace: req.Namespace(),
 		},
 		Spec: v1.OAuthAppLoginSpec{
-			CredentialContext: agent.Name,
-			ToolReference:     ref,
-			OAuthApps:         agent.Spec.Manifest.OAuthApps,
+			CredentialContext:        agent.Name,
+			ToolReference:            ref,
+			OAuthApps:                oauthApps,
+			PATSupportedIntegrations: supportedPATOAuthApps,
 		},
 	}
 
-	if err = req.Delete(oauthLogin); err != nil {
+	if err := req.Delete(oauthLogin); err != nil {
 		return err
 	}
 
-	oauthLogin, err = wait.For(req.Context(), req.Storage, oauthLogin, func(obj *v1.OAuthAppLogin) (bool, error) {
-		return obj.Status.External.Authenticated || obj.Status.External.Error != "" || obj.Status.External.URL != "", nil
+	oauthLogin, err := wait.For(req.Context(), req.Storage, oauthLogin, func(obj *v1.OAuthAppLogin) (bool, error) {
+		return obj.Status.External.Authenticated ||
+			obj.Status.External.Error != "" ||
+			obj.Status.External.URL != "" ||
+			obj.Status.External.PromptMessage != "" ||
+			len(obj.Status.External.PromptFields) != 0, nil
 	}, wait.Option{
 		Create: true,
 	})

diff --git a/pkg/api/handlers/workflows.go b/pkg/api/handlers/workflows.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net/http"
 	"strings"
 
 	"github.com/gptscript-ai/go-gptscript"
@@ -321,12 +322,16 @@ func (a *WorkflowHandler) EnsureCredentialForKnowledgeSource(req api.Context) er
 		return req.WriteCreated(resp)
 	}
 
-	credentialTools, err := v1.CredentialTools(req.Context(), req.Storage, req.Namespace(), ref)
-	if err != nil {
+	var toolReference v1.ToolReference
+	if err := req.Get(&toolReference, ref); err != nil {
 		return err
 	}
 
-	if len(credentialTools) == 0 {
+	if toolReference.Status.Tool == nil {
+		return types.NewErrHttp(http.StatusTooEarly, fmt.Sprintf("tool %q is not processed yet", ref))
+	}
+
+	if len(toolReference.Status.Tool.Credentials) == 0 {
 		// The only way to get here is if the controller hasn't set the field yet.
 		if wf.Status.AuthStatus == nil {
 			wf.Status.AuthStatus = make(map[string]types.OAuthAppLoginAuthStatus)
@@ -342,24 +347,39 @@ func (a *WorkflowHandler) EnsureCredentialForKnowledgeSource(req api.Context) er
 		return req.WriteCreated(resp)
 	}
 
+	var supportedPATOAuthApps []string
+	if toolReference.Status.Tool.Metadata["oauthPATSupported"] == "true" {
+		supportedPATOAuthApps = append(supportedPATOAuthApps, toolReference.Status.Tool.Metadata["oauth"])
+	}
+
+	var oauthApps []string
+	if req.URL.Query().Get("pat") != "true" {
+		oauthApps = append(oauthApps, toolReference.Status.Tool.Metadata["oauth"])
+	}
+
 	oauthLogin := &v1.OAuthAppLogin{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      system.OAuthAppLoginPrefix + wf.Name + ref,
 			Namespace: req.Namespace(),
 		},
 		Spec: v1.OAuthAppLoginSpec{
-			CredentialContext: wf.Name,
-			ToolReference:     ref,
-			OAuthApps:         wf.Spec.Manifest.OAuthApps,
+			CredentialContext:        wf.Name,
+			ToolReference:            ref,
+			OAuthApps:                oauthApps,
+			PATSupportedIntegrations: supportedPATOAuthApps,
 		},
 	}
 
-	if err = req.Delete(oauthLogin); err != nil {
+	if err := req.Delete(oauthLogin); err != nil {
 		return err
 	}
 
-	oauthLogin, err = wait.For(req.Context(), req.Storage, oauthLogin, func(obj *v1.OAuthAppLogin) (bool, error) {
-		return obj.Status.External.Authenticated || obj.Status.External.Error != "" || obj.Status.External.URL != "", nil
+	oauthLogin, err := wait.For(req.Context(), req.Storage, oauthLogin, func(obj *v1.OAuthAppLogin) (bool, error) {
+		return obj.Status.External.Authenticated ||
+			obj.Status.External.Error != "" ||
+			obj.Status.External.URL != "" ||
+			obj.Status.External.PromptMessage != "" ||
+			len(obj.Status.External.PromptFields) != 0, nil
 	}, wait.Option{
 		Create: true,
 	})

diff --git a/pkg/controller/handlers/oauthapp/oauthapplogin.go b/pkg/controller/handlers/oauthapp/oauthapplogin.go
@@ -97,12 +97,23 @@ outer:
 				break outer
 			}
 
-			if frame.Prompt != nil && frame.Prompt.Metadata["authURL"] != "" {
-				login.Status = v1.OAuthAppLoginStatus{
-					External: types.OAuthAppLoginAuthStatus{
-						URL:      frame.Prompt.Metadata["authURL"],
-						Required: &[]bool{true}[0],
-					},
+			if frame.Prompt != nil {
+				if frame.Prompt.Metadata["authURL"] != "" {
+					login.Status = v1.OAuthAppLoginStatus{
+						External: types.OAuthAppLoginAuthStatus{
+							URL:      frame.Prompt.Metadata["authURL"],
+							Required: &[]bool{true}[0],
+						},
+					}
+				} else {
+					login.Status = v1.OAuthAppLoginStatus{
+						External: types.OAuthAppLoginAuthStatus{
+							PromptMessage:   frame.Prompt.Message,
+							PromptFields:    frame.Prompt.Fields,
+							PromptSensitive: frame.Prompt.Sensitive,
+							Required:        &[]bool{true}[0],
+						},
+					}
 				}
 				if err = req.Client.Status().Update(req.Ctx, login); err != nil {
 					login.Status = v1.OAuthAppLoginStatus{

diff --git a/pkg/storage/apis/obot.obot.ai/v1/oauthapp.go b/pkg/storage/apis/obot.obot.ai/v1/oauthapp.go
@@ -129,9 +129,10 @@ func (o *OAuthAppLogin) DeleteRefs() []Ref {
 }
 
 type OAuthAppLoginSpec struct {
-	CredentialContext string   `json:"credentialContext,omitempty"`
-	ToolReference     string   `json:"toolReference,omitempty"`
-	OAuthApps         []string `json:"oauthApps,omitempty"`
+	CredentialContext        string   `json:"credentialContext,omitempty"`
+	ToolReference            string   `json:"toolReference,omitempty"`
+	OAuthApps                []string `json:"oauthApps,omitempty"`
+	PATSupportedIntegrations []string `json:"patSupportedIntegrations,omitempty"`
 }
 
 type OAuthAppLoginStatus struct {

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