chore: handle env vars more securely

obot-platform · Jan 14, 2025 · e7be3c6 · e7be3c6
1 parent 1af08c3
commit e7be3c6
Show file tree

Hide file tree

Showing 13 changed files with 136 additions and 86 deletions.
diff --git a/go.mod b/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/gptscript-ai/chat-completion-client v0.0.0-20241219123536-85c44096bc10
 	github.com/gptscript-ai/cmd v0.0.0-20240907001148-ffd49061124a
 	github.com/gptscript-ai/go-gptscript v0.9.6-0.20241216211344-79a66826cf82
-	github.com/gptscript-ai/gptscript v0.9.6-0.20250113160503-c4812c6df093
+	github.com/gptscript-ai/gptscript v0.9.6-0.20250114044537-9cb92bc3378a
 	github.com/liggitt/tabwriter v0.0.0-20181228230101-89fcab3d43de
 	github.com/mhale/smtpd v0.8.3
 	github.com/oauth2-proxy/oauth2-proxy/v7 v7.0.0-00010101000000-000000000000

diff --git a/go.sum b/go.sum
@@ -353,8 +353,8 @@ github.com/gptscript-ai/cmd v0.0.0-20240907001148-ffd49061124a h1:LX7AOcbBoTnUk/
 github.com/gptscript-ai/cmd v0.0.0-20240907001148-ffd49061124a/go.mod h1:DJAo1xTht1LDkNYFNydVjTHd576TC7MlpsVRl3oloVw=
 github.com/gptscript-ai/go-gptscript v0.9.6-0.20241216211344-79a66826cf82 h1:BEN268Z92gqeDc51XVvWdJWdQ47BuuWH3MUysHzilfI=
 github.com/gptscript-ai/go-gptscript v0.9.6-0.20241216211344-79a66826cf82/go.mod h1:/FVuLwhz+sIfsWUgUHWKi32qT0i6+IXlUlzs70KKt/Q=
-github.com/gptscript-ai/gptscript v0.9.6-0.20250113160503-c4812c6df093 h1:hzV7C/s5CpeybvpvabZzmDjsYZ9Eh/2RQNGx5DUZUB4=
-github.com/gptscript-ai/gptscript v0.9.6-0.20250113160503-c4812c6df093/go.mod h1:eBrKu1mmZ4tLPoHJJD1xT/Ogm5K7Oue14xk54e+yEZw=
+github.com/gptscript-ai/gptscript v0.9.6-0.20250114044537-9cb92bc3378a h1:swR4p80r0QVda+xlC2B1m9hAqR+uY0wHAzbSxjiJmWA=
+github.com/gptscript-ai/gptscript v0.9.6-0.20250114044537-9cb92bc3378a/go.mod h1:eBrKu1mmZ4tLPoHJJD1xT/Ogm5K7Oue14xk54e+yEZw=
 github.com/gptscript-ai/tui v0.0.0-20240923192013-172e51ccf1d6 h1:vkgNZVWQgbE33VD3z9WKDwuu7B/eJVVMMPM62ixfCR8=
 github.com/gptscript-ai/tui v0.0.0-20240923192013-172e51ccf1d6/go.mod h1:frrl/B+ZH3VSs3Tqk2qxEIIWTONExX3tuUa4JsVnqx4=
 github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7 h1:pdN6V1QBWetyv/0+wjACpqVH+eVULgEjkurDLq3goeM=

diff --git a/pkg/api/handlers/tools.go b/pkg/api/handlers/tools.go
@@ -28,7 +28,7 @@ func NewToolHandler(gptScript *gptscript.GPTScript, invoke *invoke.Invoker) *Too
 	}
 }
 
-var invalidEnv = regexp.MustCompile("^(OBOT|GPTSCRIPT)")
+var invalidEnv = regexp.MustCompile("^(OBOT|GPTSCRIPT|KNOW)")
 
 func setEnvMap(req api.Context, gptScript *gptscript.GPTScript, threadName, toolName string, env map[string]string) error {
 	for k := range env {

diff --git a/pkg/credstores/credstore.go b/pkg/credstores/credstore.go
@@ -17,25 +17,19 @@ type Options struct {
 
 var log = logger.Package()
 
-func Init(ctx context.Context, toolsRegistry, dsn string, opts Options) error {
+func Init(ctx context.Context, toolsRegistry, dsn string, opts Options) (string, []string, error) {
 	if err := setupKMS(ctx, opts.AWSKMSKeyARN, opts.EncryptionConfigFile); err != nil {
-		return fmt.Errorf("failed to setup kms: %w", err)
+		return "", nil, fmt.Errorf("failed to setup kms: %w", err)
 	}
 
 	switch {
 	case strings.HasPrefix(dsn, "sqlite://"):
-		if err := setupSQLite(toolsRegistry, dsn); err != nil {
-			return fmt.Errorf("failed to setup sqlite: %w", err)
-		}
+		return setupSQLite(toolsRegistry, dsn)
 	case strings.HasPrefix(dsn, "postgres://"):
-		if err := setupPostgres(toolsRegistry, dsn); err != nil {
-			return fmt.Errorf("failed to setup postgres: %w", err)
-		}
+		return setupPostgres(toolsRegistry, dsn)
 	default:
-		return fmt.Errorf("unsupported database for credentials %s", dsn)
+		return "", nil, fmt.Errorf("unsupported database for credentials %s", dsn)
 	}
-
-	return nil
 }
 
 func setupKMS(ctx context.Context, arn, configFile string) error {
@@ -75,38 +69,26 @@ func setupKMS(ctx context.Context, arn, configFile string) error {
 	return nil
 }
 
-func setupPostgres(toolRegistry, dsn string) error {
-	if err := os.Setenv("GPTSCRIPT_POSTGRES_DSN", dsn); err != nil {
-		return fmt.Errorf("failed to set GPTSCRIPT_POSTGRES_DSN: %w", err)
-	}
-
-	if err := os.Setenv("GPTSCRIPT_CREDENTIAL_STORE", toolRegistry+"/credential-stores/postgres"); err != nil {
-		return fmt.Errorf("failed to set GPTSCRIPT_CREDENTIAL_STORE: %w", err)
-	}
-
-	return nil
+func setupPostgres(toolRegistry, dsn string) (string, []string, error) {
+	return toolRegistry + "/credential-stores/postgres", []string{
+		"GPTSCRIPT_POSTGRES_DSN=" + dsn,
+	}, nil
 }
 
-func setupSQLite(toolRegistry, dsn string) error {
+func setupSQLite(toolRegistry, dsn string) (string, []string, error) {
 	dbFile, ok := strings.CutPrefix(dsn, "sqlite://file:")
 	if !ok {
-		return fmt.Errorf("invalid sqlite dsn, must start with sqlite://file: %s", dsn)
+		return "", nil, fmt.Errorf("invalid sqlite dsn, must start with sqlite://file: %s", dsn)
 	}
 	dbFile, _, _ = strings.Cut(dbFile, "?")
 
 	if !strings.HasSuffix(dbFile, ".db") {
-		return fmt.Errorf("invalid sqlite dsn, file must end in .db: %s", dsn)
+		return "", nil, fmt.Errorf("invalid sqlite dsn, file must end in .db: %s", dsn)
 	}
 
 	dbFile = strings.TrimSuffix(dbFile, ".db") + "-credentials.db"
 
-	if err := os.Setenv("GPTSCRIPT_SQLITE_FILE", dbFile); err != nil {
-		return fmt.Errorf("failed to set GPTSCRIPT_SQLITE_FILE: %w", err)
-	}
-
-	if err := os.Setenv("GPTSCRIPT_CREDENTIAL_STORE", toolRegistry+"/credential-stores/sqlite"); err != nil {
-		return fmt.Errorf("failed to set GPTSCRIPT_CREDENTIAL_STORE: %w", err)
-	}
-
-	return nil
+	return toolRegistry + "/credential-stores/sqlite", []string{
+		"GPTSCRIPT_SQLITE_FILE=" + dbFile,
+	}, nil
 }
diff --git a/pkg/invoke/system.go b/pkg/invoke/system.go
@@ -130,9 +130,15 @@ func (i *Invoker) SystemTask(ctx context.Context, thread *v1.Thread, tool, input
 		return nil, err
 	}
 
+	var credContexts []string
+	if thread != nil && thread.Namespace != "" {
+		credContexts = append(credContexts, thread.Namespace)
+	}
+	credContexts = append(opt.CredentialContextIDs, credContexts...)
+
 	return i.createRun(ctx, i.uncached, thread, tool, inputString, runOptions{
 		Env:                  opt.Env,
-		CredentialContextIDs: opt.CredentialContextIDs,
+		CredentialContextIDs: credContexts,
 		Synchronous:          true,
 		Timeout:              opt.Timeout,
 	})

diff --git a/pkg/services/config.go b/pkg/services/config.go
@@ -6,6 +6,7 @@ import (
 	"log/slog"
 	"os"
 	"path/filepath"
+	"sort"
 	"strings"
 
 	"github.com/adrg/xdg"
@@ -53,20 +54,21 @@ type (
 )
 
 type Config struct {
-	HTTPListenPort             int    `usage:"HTTP port to listen on" default:"8080" name:"http-listen-port"`
-	DevMode                    bool   `usage:"Enable development mode" default:"false" name:"dev-mode" env:"OBOT_DEV_MODE"`
-	DevUIPort                  int    `usage:"The port on localhost running the dev instance of the UI" default:"5173"`
-	AllowedOrigin              string `usage:"Allowed origin for CORS"`
-	ToolRegistry               string `usage:"The tool reference for the tool registry" default:"github.com/obot-platform/tools"`
-	WorkspaceProviderType      string `usage:"The type of workspace provider to use for non-knowledge workspaces" default:"directory" env:"OBOT_WORKSPACE_PROVIDER_TYPE"`
-	WorkspaceTool              string `usage:"The tool reference for the workspace provider" default:"github.com/gptscript-ai/workspace-provider"`
-	DatasetsTool               string `usage:"The tool reference for the dataset provider" default:"github.com/gptscript-ai/datasets"`
-	HelperModel                string `usage:"The model used to generate names and descriptions" default:"gpt-4o-mini"`
-	AWSKMSKeyARN               string `usage:"The ARN of the AWS KMS key to use for encrypting credential storage" env:"OBOT_AWS_KMS_KEY_ARN" name:"aws-kms-key-arn"`
-	EncryptionConfigFile       string `usage:"The path to the encryption configuration file" default:"./encryption.yaml"`
-	KnowledgeSetIngestionLimit int    `usage:"The maximum number of files to ingest into a knowledge set" default:"3000" env:"OBOT_KNOWLEDGESET_INGESTION_LIMIT" name:"knowledge-set-ingestion-limit"`
-	EmailServerName            string `usage:"The name of the email server to display for email receivers"`
-	Docker                     bool   `usage:"Enable Docker support" default:"false" env:"OBOT_DOCKER"`
+	HTTPListenPort             int      `usage:"HTTP port to listen on" default:"8080" name:"http-listen-port"`
+	DevMode                    bool     `usage:"Enable development mode" default:"false" name:"dev-mode" env:"OBOT_DEV_MODE"`
+	DevUIPort                  int      `usage:"The port on localhost running the dev instance of the UI" default:"5173"`
+	AllowedOrigin              string   `usage:"Allowed origin for CORS"`
+	ToolRegistry               string   `usage:"The tool reference for the tool registry" default:"github.com/obot-platform/tools"`
+	WorkspaceProviderType      string   `usage:"The type of workspace provider to use for non-knowledge workspaces" default:"directory" env:"OBOT_WORKSPACE_PROVIDER_TYPE"`
+	WorkspaceTool              string   `usage:"The tool reference for the workspace provider" default:"github.com/gptscript-ai/workspace-provider"`
+	DatasetsTool               string   `usage:"The tool reference for the dataset provider" default:"github.com/gptscript-ai/datasets"`
+	HelperModel                string   `usage:"The model used to generate names and descriptions" default:"gpt-4o-mini"`
+	AWSKMSKeyARN               string   `usage:"The ARN of the AWS KMS key to use for encrypting credential storage" env:"OBOT_AWS_KMS_KEY_ARN" name:"aws-kms-key-arn"`
+	EncryptionConfigFile       string   `usage:"The path to the encryption configuration file" default:"./encryption.yaml"`
+	KnowledgeSetIngestionLimit int      `usage:"The maximum number of files to ingest into a knowledge set" default:"3000" env:"OBOT_KNOWLEDGESET_INGESTION_LIMIT" name:"knowledge-set-ingestion-limit"`
+	EmailServerName            string   `usage:"The name of the email server to display for email receivers"`
+	Docker                     bool     `usage:"Enable Docker support" default:"false" env:"OBOT_DOCKER"`
+	EnvKeys                    []string `usage:"The environment keys to pass through to the GPTScript server" env:"OBOT_ENV_KEYS"`
 
 	AuthConfig
 	GatewayConfig
@@ -101,7 +103,39 @@ const (
 	defaultToolsRegistry = "github.com/obot-platform/tools"
 )
 
-func newGPTScript(ctx context.Context, workspaceTool, datasetsTool, toolsRegistry string) (*gptscript.GPTScript, error) {
+var requiredEnvs = []string{
+	// Standard system stuff
+	"PATH", "HOME", "USER", "PWD",
+	// Embedded env vars
+	"OBOT_BIN", "GPTSCRIPT_BIN", "GPTSCRIPT_EMBEDDED",
+	// XDG stuff
+	"XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_CACHE_HOME"}
+
+func copyKeys(envs []string) []string {
+	seen := make(map[string]struct{})
+	newEnvs := make([]string, len(envs))
+
+	for _, env := range append(envs, requiredEnvs...) {
+		if env == "*" {
+			return os.Environ()
+		}
+		if _, ok := seen[env]; ok {
+			continue
+		}
+		v := os.Getenv(env)
+		if v == "" {
+			continue
+		}
+		seen[env] = struct{}{}
+		newEnvs = append(newEnvs, fmt.Sprintf("%s=%s", env, os.Getenv(env)))
+	}
+
+	sort.Strings(newEnvs)
+	return newEnvs
+}
+
+func newGPTScript(ctx context.Context, workspaceTool, datasetsTool, toolsRegistry string, envPassThrough []string,
+	credStore string, credStoreEnv []string) (*gptscript.GPTScript, error) {
 	if datasetsTool != defaultDatasetsTool {
 		loader.Remap[defaultDatasetsTool] = datasetsTool
 	}
@@ -123,13 +157,16 @@ func newGPTScript(ctx context.Context, workspaceTool, datasetsTool, toolsRegistr
 	}
 	url, err := sdkserver.EmbeddedStart(ctx, sdkserver.Options{
 		Options: gptscriptai.Options{
+			Env: copyKeys(envPassThrough),
 			Cache: cache.Options{
 				CacheDir: os.Getenv("GPTSCRIPT_CACHE_DIR"),
 			},
 			Runner: runner.Options{
 				CredentialOverrides: credOverrides,
 			},
-			SystemToolsDir: os.Getenv("GPTSCRIPT_SYSTEM_TOOLS_DIR"),
+			SystemToolsDir:     os.Getenv("GPTSCRIPT_SYSTEM_TOOLS_DIR"),
+			CredentialStore:    credStore,
+			CredentialToolsEnv: append(copyKeys(envPassThrough), credStoreEnv...),
 		},
 		DatasetTool:   datasetsTool,
 		WorkspaceTool: workspaceTool,
@@ -149,6 +186,7 @@ func newGPTScript(ctx context.Context, workspaceTool, datasetsTool, toolsRegistr
 	}
 
 	return gptscript.NewGPTScript(gptscript.GlobalOptions{
+		Env:           copyKeys(envPassThrough),
 		URL:           url,
 		WorkspaceTool: workspaceTool,
 		DatasetTool:   datasetsTool,
@@ -163,15 +201,12 @@ func New(ctx context.Context, config Config) (*Services, error) {
 	// Just a common mistake where you put the wrong prefix for the DSN. This seems to be inconsistent across things
 	// that use postgres
 	config.DSN = strings.Replace(config.DSN, "postgresql://", "postgres://", 1)
-	if strings.HasPrefix(config.DSN, "postgres://") {
-		_ = os.Setenv("KNOW_VECTOR_DSN", strings.Replace(config.DSN, "postgres://", "pgvector://", 1))
-		_ = os.Setenv("KNOW_INDEX_DSN", config.DSN)
-	}
 
-	if err := credstores.Init(ctx, config.ToolRegistry, config.DSN, credstores.Options{
+	credStore, credStoreEnv, err := credstores.Init(ctx, config.ToolRegistry, config.DSN, credstores.Options{
 		AWSKMSKeyARN:         config.AWSKMSKeyARN,
 		EncryptionConfigFile: config.EncryptionConfigFile,
-	}); err != nil {
+	})
+	if err != nil {
 		return nil, err
 	}
 
@@ -205,11 +240,30 @@ func New(ctx context.Context, config Config) (*Services, error) {
 		config.UIHostname = "https://" + config.UIHostname
 	}
 
-	c, err := newGPTScript(ctx, config.WorkspaceTool, config.DatasetsTool, config.ToolRegistry)
+	c, err := newGPTScript(ctx, config.WorkspaceTool, config.DatasetsTool, config.ToolRegistry, config.EnvKeys,
+		credStore, credStoreEnv)
 	if err != nil {
 		return nil, err
 	}
 
+	if strings.HasPrefix(config.DSN, "postgres://") {
+		if err := c.CreateCredential(ctx, gptscript.Credential{
+			Context:  system.DefaultNamespace,
+			ToolName: system.KnowledgeCredID,
+			Type:     gptscript.CredentialTypeTool,
+			Env: map[string]string{
+				"KNOW_VECTOR_DSN": strings.Replace(config.DSN, "postgres://", "pgvector://", 1),
+				"KNOW_INDEX_DSN":  config.DSN,
+			},
+		}); err != nil {
+			return nil, err
+		}
+	} else {
+		if err := c.DeleteCredential(ctx, system.DefaultNamespace, system.KnowledgeCredID); err != nil {
+			return nil, err
+		}
+	}
+
 	r, err := baaah.NewRouter("obot-controller", &baaah.Options{
 		DefaultRESTConfig: restConfig,
 		Scheme:            scheme.Scheme,

diff --git a/pkg/system/tools.go b/pkg/system/tools.go
@@ -14,6 +14,7 @@ const (
 	ShellTool               = "shell"
 	DockerShellIDTool       = "docker-shell-id"
 	ExistingCredTool        = "existing-credential"
+	KnowledgeCredID         = "knowledge"
 
 	DefaultNamespace = "default"
 

diff --git a/ui/user/src/lib/components/Thread.svelte b/ui/user/src/lib/components/Thread.svelte
@@ -22,8 +22,8 @@
 		if (!a && $currentAssistant.id === assistant) {
 			a = $currentAssistant;
 		}
-		return a
-	})
+		return a;
+	});
 
 	$effect(() => {
 		if (!assistant || thread) {
@@ -64,15 +64,18 @@
 			<div in:fade|global class="flex flex-col gap-8">
 				{#if messages.messages.length < 7}
 					<div class="message-content self-center">
-						{#if current?.introductionMessage }
+						{#if current?.introductionMessage}
 							{@html toHTMLFromMarkdown(current.introductionMessage)}
 						{/if}
 					</div>
 					<div class="flex gap-2 self-center">
 						{#each current?.starterMessages ?? [] as msg}
-							<button class="border-2 border-blue rounded-3xl p-5" onclick={() => {
-								thread?.invoke(msg);
-							}}>
+							<button
+								class="rounded-3xl border-2 border-blue p-5"
+								onclick={() => {
+									thread?.invoke(msg);
+								}}
+							>
 								{msg}
 							</button>
 						{/each}

diff --git a/ui/user/src/lib/components/navbar/KnowledgeFile.svelte b/ui/user/src/lib/components/navbar/KnowledgeFile.svelte
@@ -29,7 +29,7 @@
 				{file.error ? file.error : 'Failed'}
 			</div>
 		{:else if file.state === 'pending' || file.state === 'ingesting'}
-			<Loading />
+			<Loading class="mx-1.5" />
 		{/if}
 	</button>
 	<button

diff --git a/ui/user/src/lib/components/navbar/KnowledgeUpload.svelte b/ui/user/src/lib/components/navbar/KnowledgeUpload.svelte
@@ -4,7 +4,7 @@
 	import type { KnowledgeFile } from '$lib/services';
 	import Loading from '$lib/icons/Loading.svelte';
 	import Error from '$lib/components/Error.svelte';
-	import { currentAssistant } from '$lib/stores';
+	import { currentAssistant, knowledgeFiles } from '$lib/stores';
 
 	interface Props {
 		onUpload?: () => void | Promise<void>;
@@ -15,6 +15,18 @@
 	let files = $state<FileList>();
 	let uploadInProgress = $state<Promise<KnowledgeFile>>();
 
+	function reloadFiles() {
+		ChatService.listKnowledgeFiles($currentAssistant.id).then((files) => {
+			knowledgeFiles.set(files);
+			const pending = files.items.find(
+				(file) => file.state === 'pending' || file.state === 'ingesting'
+			);
+			if (pending) {
+				setTimeout(reloadFiles, 2000);
+			}
+		});
+	}
+
 	$effect(() => {
 		if (!files?.length) {
 			return;
@@ -30,6 +42,7 @@
 			})
 			.finally(() => {
 				uploadInProgress = undefined;
+				setTimeout(reloadFiles, 1000);
 			});
 
 		files = undefined;

diff --git a/ui/user/src/lib/components/terminal/Env.svelte b/ui/user/src/lib/components/terminal/Env.svelte
@@ -1,5 +1,4 @@
 <script lang="ts">
-	import { onMount } from 'svelte';
 	import Env from '$lib/components/tool/Env.svelte';
 	import { ChatService } from '$lib/services';
 	import { currentAssistant } from '$lib/stores';
@@ -39,10 +38,6 @@
 		}
 		dialog.close();
 	}
-
-	onMount(() => {
-		show();
-	});
 </script>
 
 <dialog