fix: use custom oauth configurations for knowledge syncing (#597)

Signed-off-by: Donnie Adams <[email protected]>

pkg/api/handlers/agent.go

@@ -292,7 +292,7 @@ func (a *AgentHandler) CreateKnowledgeSource(req api.Context) error {
 
 	var input types.KnowledgeSourceManifest
 	if err := req.Read(&input); err != nil {
-		return types.NewErrBadRequest("failed to decode request body: %w", err)
+		return types.NewErrBadRequest("failed to decode request body: %v", err)
 	}
 
 	if err := input.Validate(); err != nil {
@@ -312,7 +312,7 @@ func (a *AgentHandler) CreateKnowledgeSource(req api.Context) error {
 	}
 
 	if err := req.Create(&source); err != nil {
-		return types.NewErrBadRequest("failed to create RemoteKnowledgeSource: %w", err)
+		return types.NewErrBadRequest("failed to create RemoteKnowledgeSource: %v", err)
 	}
 
 	return req.Write(convertKnowledgeSource(agentName, source))
@@ -326,7 +326,7 @@ func (a *AgentHandler) UpdateKnowledgeSource(req api.Context) error {
 
 	var manifest types.KnowledgeSourceManifest
 	if err := req.Read(&manifest); err != nil {
-		return types.NewErrBadRequest("failed to decode request body: %w", err)
+		return types.NewErrBadRequest("failed to decode request body: %v", err)
 	}
 
 	if err := manifest.Validate(); err != nil {
@@ -535,6 +535,7 @@ func (a *AgentHandler) EnsureCredentialForKnowledgeSource(req api.Context) error
 		Spec: v1.OAuthAppLoginSpec{
 			CredentialContext: agent.Name,
 			ToolReference:     ref,
+			OAuthApps:         agent.Spec.Manifest.OAuthApps,
 		},
 	}
 
@@ -566,7 +567,7 @@ func (a *AgentHandler) Script(req api.Context) error {
 		agent v1.Agent
 	)
 	if err := req.Get(&agent, id); err != nil {
-		return types.NewErrBadRequest("failed to get agent with id %s: %w", id, err)
+		return types.NewErrBadRequest("failed to get agent with id %s: %v", id, err)
 	}
 
 	tools, extraEnv, err := render.Agent(req.Context(), req.Storage, &agent, a.serverURL, render.AgentOptions{})

pkg/api/handlers/assistants.go

@@ -346,7 +346,7 @@ func (a *AssistantHandler) DeleteKnowledge(req api.Context) error {
 	}
 
 	if len(thread.Status.KnowledgeSetNames) == 0 {
-		return types.NewErrHttp(http.StatusTooEarly, fmt.Sprintf("knowledge set is not created yet"))
+		return types.NewErrHttp(http.StatusTooEarly, "knowledge set is not created yet")
 	}
 
 	return deleteKnowledge(req, req.PathValue("file"), thread.Status.KnowledgeSetNames[0])

pkg/api/handlers/webhooks.go

@@ -283,7 +283,7 @@ func validateSecretHeader(secret string, body []byte, values []string) error {
 func validateManifest(req api.Context, manifest types.WebhookManifest) error {
 	// Ensure that the WorkflowID is set and the workflow exists
 	if manifest.WorkflowID == "" {
-		return apierrors.NewBadRequest(fmt.Sprintf("webhook manifest must have a workflow name"))
+		return apierrors.NewBadRequest("webhook manifest must have a workflow name")
 	}
 
 	var workflow v1.Workflow
@@ -297,7 +297,7 @@ func validateManifest(req api.Context, manifest types.WebhookManifest) error {
 	}
 
 	if (manifest.ValidationHeader != "") != (manifest.Secret != "") {
-		return apierrors.NewBadRequest(fmt.Sprintf("webhook must have secret and header set together"))
+		return apierrors.NewBadRequest("webhook must have secret and header set together")
 	}
 
 	return nil

pkg/api/handlers/workflows.go

@@ -209,6 +209,7 @@ func (a *WorkflowHandler) EnsureCredentialForKnowledgeSource(req api.Context) er
 		Spec: v1.OAuthAppLoginSpec{
 			CredentialContext: wf.Name,
 			ToolReference:     ref,
+			OAuthApps:         wf.Spec.Manifest.OAuthApps,
 		},
 	}
 

pkg/controller/handlers/oauthapp/oauthapplogin.go

@@ -8,6 +8,7 @@ import (
 	"github.com/otto8-ai/nah/pkg/router"
 	"github.com/otto8-ai/otto8/apiclient/types"
 	"github.com/otto8-ai/otto8/pkg/invoke"
+	"github.com/otto8-ai/otto8/pkg/render"
 	v1 "github.com/otto8-ai/otto8/pkg/storage/apis/otto.otto8.ai/v1"
 	"github.com/otto8-ai/otto8/pkg/system"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -16,12 +17,14 @@ import (
 )
 
 type LoginHandler struct {
-	invoker *invoke.Invoker
+	invoker   *invoke.Invoker
+	serverURL string
 }
 
-func NewLogin(invoker *invoke.Invoker) *LoginHandler {
+func NewLogin(invoker *invoke.Invoker, serverURL string) *LoginHandler {
 	return &LoginHandler{
-		invoker: invoker,
+		invoker:   invoker,
+		serverURL: serverURL,
 	}
 }
 
@@ -50,13 +53,19 @@ 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)
+	if err != nil {
+		return err
+	}
+
 	task, err := h.invoker.SystemTask(req.Ctx, &thread, []gptscript.ToolDef{
 		{
 			Credentials:  []string{credentialTool},
 			Instructions: "#!sys.echo DONE",
 		},
 	}, "", invoke.SystemTaskOptions{
 		CredentialContextIDs: []string{login.Spec.CredentialContext},
+		Env:                  oauthAppEnv,
 	})
 	if err != nil {
 		return err

pkg/controller/routes.go

@@ -34,7 +34,7 @@ func (c *Controller) setupRoutes() error {
 	runs := runs.New(c.services.Invoker)
 	webHooks := webhook.New()
 	cronJobs := cronjob.New()
-	oauthLogins := oauthapp.NewLogin(c.services.Invoker)
+	oauthLogins := oauthapp.NewLogin(c.services.Invoker, c.services.ServerURL)
 
 	// Runs
 	root.Type(&v1.Run{}).FinalizeFunc(v1.RunFinalizer, runs.DeleteRunState)

pkg/invoke/system.go

@@ -9,11 +9,13 @@ import (
 
 type SystemTaskOptions struct {
 	CredentialContextIDs []string
+	Env                  []string
 }
 
 func complete(opts []SystemTaskOptions) (result SystemTaskOptions) {
 	for _, opt := range opts {
 		result.CredentialContextIDs = append(result.CredentialContextIDs, opt.CredentialContextIDs...)
+		result.Env = append(result.Env, opt.Env...)
 	}
 	return
 }
@@ -42,6 +44,7 @@ func (i *Invoker) SystemTask(ctx context.Context, thread *v1.Thread, tool, input
 	}
 
 	return i.createRun(ctx, i.uncached, thread, tool, inputString, runOptions{
+		Env:                  opt.Env,
 		CredentialContextIDs: opt.CredentialContextIDs,
 		Synchronous:          true,
 	})

pkg/render/render.go

@@ -78,7 +78,7 @@ func Agent(ctx context.Context, db kclient.Client, agent *v1.Agent, oauthServerU
 		return nil, nil, err
 	}
 
-	if oauthEnv, err := setupOAuthApps(ctx, db, agent, oauthServerURL); err != nil {
+	if oauthEnv, err := OAuthAppEnv(ctx, db, agent.Spec.Manifest.OAuthApps, agent.Namespace, oauthServerURL); err != nil {
 		return nil, nil, err
 	} else {
 		extraEnv = append(extraEnv, oauthEnv...)
@@ -87,8 +87,8 @@ func Agent(ctx context.Context, db kclient.Client, agent *v1.Agent, oauthServerU
 	return append([]gptscript.ToolDef{mainTool}, otherTools...), extraEnv, nil
 }
 
-func setupOAuthApps(ctx context.Context, db kclient.Client, agent *v1.Agent, serverURL string) (extraEnv []string, _ error) {
-	apps, err := oauthAppsByName(ctx, db, agent.Namespace)
+func OAuthAppEnv(ctx context.Context, db kclient.Client, oauthAppNames []string, namespace, serverURL string) (extraEnv []string, _ error) {
+	apps, err := oauthAppsByName(ctx, db, namespace)
 	if err != nil {
 		return nil, err
 	}
@@ -102,7 +102,7 @@ func setupOAuthApps(ctx context.Context, db kclient.Client, agent *v1.Agent, ser
 		activeIntegrations[app.Spec.Manifest.Integration] = app
 	}
 
-	for _, appRef := range agent.Spec.Manifest.OAuthApps {
+	for _, appRef := range oauthAppNames {
 		app, ok := apps[appRef]
 		if !ok {
 			return nil, fmt.Errorf("oauth app %s not found", appRef)

pkg/storage/apis/otto.otto8.ai/v1/oauthapp.go

@@ -158,8 +158,9 @@ func (o *OAuthAppLogin) DeleteRefs() []Ref {
 }
 
 type OAuthAppLoginSpec struct {
-	CredentialContext string `json:"credentialContext,omitempty"`
-	ToolReference     string `json:"toolReference,omitempty"`
+	CredentialContext string   `json:"credentialContext,omitempty"`
+	ToolReference     string   `json:"toolReference,omitempty"`
+	OAuthApps         []string `json:"oauthApps,omitempty"`
 }
 
 type OAuthAppLoginStatus struct {