Skip to content

Commit

Permalink
feat: add support for OAuth PATs
Browse files Browse the repository at this point in the history
Instead of configuring and using OAuth, an app can specify that it
supports using personal access tokens. If this is the case, then Obot
will pass an extra environment variable to the oauth credential tool to
indicate which integrations support tokens. If the oauth2 credential
tool should prompt the user for a token instead of using OAuth, then
Obot will not pass the environment variables that feed the URLs to the
tool.

A side effect of this change is that OAuth apps no longer default to
global.

Signed-off-by: Donnie Adams <[email protected]>
  • Loading branch information
thedadams committed Jan 22, 2025
1 parent bd46fac commit 76db187
Show file tree
Hide file tree
Showing 12 changed files with 83 additions and 39 deletions.
2 changes: 1 addition & 1 deletion apiclient/types/oauthapp.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ type OAuthAppManifest struct {
// This field is required, it correlates to the integration name in the gptscript oauth cred tool
Integration string `json:"integration,omitempty"`
// Global indicates if the OAuth app is globally applied to all agents.
Global *bool `json:"global,omitempty"`
Global bool `json:"global,omitempty"`
// This field is only used by Salesforce
InstanceURL string `json:"instanceURL,omitempty"`
}
Expand Down
5 changes: 0 additions & 5 deletions apiclient/types/zz_generated.deepcopy.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions pkg/api/handlers/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -918,7 +918,7 @@ func runAuthForAgent(ctx context.Context, c kclient.WithWatch, invoker *invoke.I

var toolRef v1.ToolReference
for _, tool := range tools {
if strings.ContainsAny(tool, "./") {
if render.IsExternalTool(tool) {
prg, err := gClient.LoadFile(ctx, tool)
if err != nil {
return nil, err
Expand Down Expand Up @@ -965,7 +965,7 @@ func removeToolCredentials(ctx context.Context, client kclient.Client, gClient *
credentialNames []string
)
for _, tool := range tools {
if strings.ContainsAny(tool, "./") {
if render.IsExternalTool(tool) {
prg, err := gClient.LoadFile(ctx, tool)
if err != nil {
errs = append(errs, err)
Expand Down
2 changes: 1 addition & 1 deletion pkg/controller/handlers/oauthapp/oauthapplogin.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ func (h *LoginHandler) RunTool(req router.Request, _ router.Response) error {
return err
}

oauthAppEnv, err := render.OAuthAppEnv(req.Ctx, req.Client, login.Spec.OAuthApps, login.Namespace, h.serverURL)
oauthAppEnv, err := render.OAuthAppEnv(req.Ctx, req.Client, login.Spec.OAuthApps, login.Namespace, h.serverURL, login.Spec.PATSupportedIntegrations)
if err != nil {
return err
}
Expand Down
4 changes: 2 additions & 2 deletions pkg/controller/handlers/toolinfo/toolinfo.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,12 @@ package toolinfo
import (
"context"
"fmt"
"strings"

"github.com/gptscript-ai/go-gptscript"
"github.com/obot-platform/nah/pkg/router"
"github.com/obot-platform/obot/apiclient/types"
"github.com/obot-platform/obot/pkg/controller/creds"
"github.com/obot-platform/obot/pkg/render"
v1 "github.com/obot-platform/obot/pkg/storage/apis/obot.obot.ai/v1"
apierror "k8s.io/apimachinery/pkg/api/errors"
"k8s.io/apimachinery/pkg/util/sets"
Expand Down Expand Up @@ -57,7 +57,7 @@ func (h *Handler) SetToolInfoStatus(req router.Request, resp router.Response) (e
credNames []string
)
for _, tool := range tools {
if strings.ContainsAny(tool, "/.") {
if render.IsExternalTool(tool) {
credNames, err = h.credentialNamesForNonToolReferences(req.Ctx, tool)
if err != nil {
return err
Expand Down
2 changes: 1 addition & 1 deletion pkg/gateway/types/oauth_apps.go
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ func MergeOAuthAppManifests(r, other types.OAuthAppManifest) types.OAuthAppManif
if other.OptionalScope != "" {
retVal.OptionalScope = other.OptionalScope
}
if other.Global != nil {
if other.Global {
retVal.Global = other.Global
}

Expand Down
39 changes: 29 additions & 10 deletions pkg/render/render.go
Original file line number Diff line number Diff line change
Expand Up @@ -74,11 +74,11 @@ func Agent(ctx context.Context, db kclient.Client, agent *v1.Agent, oauthServerU
}

if opts.Thread != nil {
for _, tool := range opts.Thread.Spec.Manifest.Tools {
if !added && tool == knowledgeToolName {
for _, t := range opts.Thread.Spec.Manifest.Tools {
if !added && t == knowledgeToolName {
continue
}
name, tools, err := Tool(ctx, db, agent.Namespace, tool)
name, _, tools, err := tool(ctx, db, agent.Namespace, t)
if err != nil {
return nil, nil, err
}
Expand All @@ -99,18 +99,33 @@ func Agent(ctx context.Context, db kclient.Client, agent *v1.Agent, oauthServerU
}
}

for _, tool := range agent.Spec.Manifest.Tools {
if !added && tool == knowledgeToolName {
var patSupportedIntegrations []string
for _, t := range agent.Spec.Manifest.Tools {
if !added && t == knowledgeToolName {
continue
}
name, tools, err := Tool(ctx, db, agent.Namespace, tool)
name, metadata, tools, err := tool(ctx, db, agent.Namespace, t)
if err != nil {
return nil, nil, err
}
if name != "" {
mainTool.Tools = append(mainTool.Tools, name)
}
otherTools = append(otherTools, tools...)

if metadata["oauthPATSupported"] == "true" {
if integration := metadata["oauth"]; integration != "" {
patSupportedIntegrations = append(patSupportedIntegrations, integration)
}
}

for _, t := range tools {
if t.MetaData["oauthPATSupported"] == "true" {
if integration := t.MetaData["oauth"]; integration != "" {
patSupportedIntegrations = append(patSupportedIntegrations, integration)
}
}
otherTools = append(otherTools, t)
}
}

for _, tool := range agent.Spec.SystemTools {
Expand All @@ -134,7 +149,7 @@ func Agent(ctx context.Context, db kclient.Client, agent *v1.Agent, oauthServerU
return nil, nil, err
}

oauthEnv, err := OAuthAppEnv(ctx, db, agent.Spec.Manifest.OAuthApps, agent.Namespace, oauthServerURL)
oauthEnv, err := OAuthAppEnv(ctx, db, agent.Spec.Manifest.OAuthApps, agent.Namespace, oauthServerURL, patSupportedIntegrations)
if err != nil {
return nil, nil, err
}
Expand All @@ -144,7 +159,7 @@ func Agent(ctx context.Context, db kclient.Client, agent *v1.Agent, oauthServerU
return append([]gptscript.ToolDef{mainTool}, otherTools...), extraEnv, nil
}

func OAuthAppEnv(ctx context.Context, db kclient.Client, oauthAppNames []string, namespace, serverURL string) (extraEnv []string, _ error) {
func OAuthAppEnv(ctx context.Context, db kclient.Client, oauthAppNames []string, namespace, serverURL string, patIntegrations []string) (extraEnv []string, _ error) {
apps, err := oauthAppsByName(ctx, db, namespace)
if err != nil {
return nil, err
Expand All @@ -153,7 +168,7 @@ func OAuthAppEnv(ctx context.Context, db kclient.Client, oauthAppNames []string,
activeIntegrations := map[string]v1.OAuthApp{}
for _, name := range slices.Sorted(maps.Keys(apps)) {
app := apps[name]
if app.Spec.Manifest.Global == nil || !*app.Spec.Manifest.Global || app.Spec.Manifest.ClientID == "" || app.Spec.Manifest.ClientSecret == "" || app.Spec.Manifest.Integration == "" {
if !app.Spec.Manifest.Global || app.Spec.Manifest.ClientID == "" || app.Spec.Manifest.ClientSecret == "" || app.Spec.Manifest.Integration == "" {
continue
}
activeIntegrations[app.Spec.Manifest.Integration] = app
Expand Down Expand Up @@ -184,6 +199,10 @@ func OAuthAppEnv(ctx context.Context, db kclient.Client, oauthAppNames []string,
fmt.Sprintf("GPTSCRIPT_OAUTH_%s_TOKEN_URL=%s", integrationEnv, v1.OAuthAppGetTokenURL(serverURL)))
}

if len(patIntegrations) > 0 {
extraEnv = append(extraEnv, fmt.Sprintf("GPTSCRIPT_OAUTH_PAT_INTEGRATIONS=%s", strings.Join(patIntegrations, ",")))
}

return extraEnv, nil
}

Expand Down
14 changes: 7 additions & 7 deletions pkg/render/tool.go
Original file line number Diff line number Diff line change
Expand Up @@ -93,25 +93,25 @@ END INSTRUCTIONS: TOOL %q`, tool.Spec.Manifest.Name, tool.Spec.Manifest.Context,
return toolDefs, nil
}

func Tool(ctx context.Context, c client.Client, ns, name string) (_ string, toolDefs []gptscript.ToolDef, _ error) {
func tool(ctx context.Context, c client.Client, ns, name string) (string, map[string]string, []gptscript.ToolDef, error) {
if !system.IsToolID(name) {
name, err := ResolveToolReference(ctx, c, types.ToolReferenceTypeTool, ns, name)
return name, nil, err
name, metadata, err := resolveToolReferenceWithMetadata(ctx, c, types.ToolReferenceTypeTool, ns, name)
return name, metadata, nil, err
}

var tool v1.Tool
if err := c.Get(ctx, router.Key(ns, name), &tool); err != nil {
return name, nil, err
return name, nil, nil, err
}

toolDefs, err := CustomTool(ctx, c, tool)
if err != nil {
return "", nil, err
return "", nil, nil, err
}

if len(toolDefs) == 0 {
return "", toolDefs, nil
return "", nil, toolDefs, nil
}

return toolDefs[0].Name, toolDefs, nil
return toolDefs[0].Name, nil, toolDefs, nil
}
24 changes: 17 additions & 7 deletions pkg/render/workflow.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,26 +29,36 @@ func IsExternalTool(tool string) bool {
}

func ResolveToolReference(ctx context.Context, c kclient.Client, toolRefType types.ToolReferenceType, ns, name string) (string, error) {
name, _, err := resolveToolReferenceWithMetadata(ctx, c, toolRefType, ns, name)
return name, err
}

func resolveToolReferenceWithMetadata(ctx context.Context, c kclient.Client, toolRefType types.ToolReferenceType, ns, name string) (string, map[string]string, error) {
if IsExternalTool(name) {
return name, nil
return name, nil, nil
}

var tool v1.ToolReference
if err := c.Get(ctx, router.Key(ns, name), &tool); apierror.IsNotFound(err) {
return name, nil
return name, nil, nil
} else if err != nil {
return "", err
return "", nil, err
}

var metadata map[string]string
if tool.Status.Tool != nil {
metadata = tool.Status.Tool.Metadata
}
if toolRefType != "" && tool.Spec.Type != toolRefType {
return name, fmt.Errorf("tool reference %s is not of type %s", name, toolRefType)
return name, metadata, fmt.Errorf("tool reference %s is not of type %s", name, toolRefType)
}
if tool.Status.Reference == "" {
return "", fmt.Errorf("tool reference %s has no reference", name)
return "", nil, fmt.Errorf("tool reference %s has no reference", name)
}
if toolRefType == types.ToolReferenceTypeTool {
return fmt.Sprintf("%s as %s", tool.Status.Reference, name), nil
return fmt.Sprintf("%s as %s", tool.Status.Reference, name), metadata, nil
}
return tool.Status.Reference, nil
return tool.Status.Reference, metadata, nil
}

func Workflow(ctx context.Context, c kclient.Client, wf *v1.Workflow, opts WorkflowOptions) (*v1.Agent, error) {
Expand Down
7 changes: 4 additions & 3 deletions pkg/storage/apis/obot.obot.ai/v1/oauthapp.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
5 changes: 5 additions & 0 deletions pkg/storage/apis/obot.obot.ai/v1/zz_generated.deepcopy.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

14 changes: 14 additions & 0 deletions pkg/storage/openapi/generated/openapi_generated.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 76db187

Please sign in to comment.