chore: compress env vars that are too large

Author: ibuildthecloud

pkg/cli/getenv.go

@@ -0,0 +1,60 @@
+package cli
+
+import (
+	"bytes"
+	"compress/gzip"
+	"encoding/base64"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/spf13/cobra"
+)
+
+type Getenv struct {
+}
+
+func (e *Getenv) Customize(cmd *cobra.Command) {
+	cmd.Use = "getenv [flags] KEY [DEFAULT]"
+	cmd.Short = "Looks up an environment variable for use in GPTScript tools"
+	cmd.Args = cobra.RangeArgs(1, 2)
+}
+
+func (e *Getenv) Run(_ *cobra.Command, args []string) error {
+	var (
+		key = args[0]
+		def string
+	)
+	if len(args) > 1 {
+		def = args[1]
+	}
+	value := getEnv(key, def)
+	fmt.Print(value)
+	return nil
+}
+
+func getEnv(key, def string) string {
+	v := os.Getenv(key)
+	if v == "" {
+		return def
+	}
+
+	if strings.HasPrefix(v, `{"_gz":"`) && strings.HasSuffix(v, `"}`) {
+		data, err := base64.StdEncoding.DecodeString(v[8 : len(v)-2])
+		if err != nil {
+			return v
+		}
+		gz, err := gzip.NewReader(bytes.NewBuffer(data))
+		if err != nil {
+			return v
+		}
+		strBytes, err := io.ReadAll(gz)
+		if err != nil {
+			return v
+		}
+		return string(strBytes)
+	}
+
+	return v
+}

pkg/cli/getenv_test.go

@@ -0,0 +1,72 @@
+package cli
+
+import (
+	"bytes"
+	"compress/gzip"
+	"encoding/base64"
+	"os"
+	"testing"
+)
+
+func TestGetEnv(t *testing.T) {
+	// Cleaning up
+	defer func(currentEnvValue string) {
+		os.Setenv("testKey", currentEnvValue)
+	}(os.Getenv("testKey"))
+
+	// Tests
+	testCases := []struct {
+		name           string
+		key            string
+		def            string
+		envValue       string
+		expectedResult string
+	}{
+		{
+			name:           "NoValueUseDefault",
+			key:            "testKey",
+			def:            "defaultValue",
+			envValue:       "",
+			expectedResult: "defaultValue",
+		},
+		{
+			name:           "ValueExistsNoCompress",
+			key:            "testKey",
+			def:            "defaultValue",
+			envValue:       "testValue",
+			expectedResult: "testValue",
+		},
+		{
+			name:     "ValueExistsCompressed",
+			key:      "testKey",
+			def:      "defaultValue",
+			envValue: `{"_gz":"H4sIAEosrGYC/ytJLS5RKEvMKU0FACtB3ewKAAAA"}`,
+
+			expectedResult: "test value",
+		},
+	}
+
+	for _, test := range testCases {
+		t.Run(test.name, func(t *testing.T) {
+			os.Setenv(test.key, test.envValue)
+
+			result := getEnv(test.key, test.def)
+
+			if result != test.expectedResult {
+				t.Errorf("expected: %s, got: %s", test.expectedResult, result)
+			}
+		})
+	}
+}
+
+// Helper Function to Get Compressed String for Testing
+func getCompressedString(s string) string {
+	var buf bytes.Buffer
+	gz := gzip.NewWriter(&buf)
+
+	_, _ = gz.Write([]byte(s))
+	_ = gz.Close()
+
+	encodedStr := base64.StdEncoding.EncodeToString(buf.Bytes())
+	return `{"_gz":"` + encodedStr + `"}`
+}

pkg/cli/gptscript.go

@@ -85,6 +85,7 @@ func New() *cobra.Command {
 		&Credential{root: root},
 		&Parse{},
 		&Fmt{},
+		&Getenv{},
 		&SDKServer{
 			GPTScript: root,
 		},

pkg/engine/cmd.go

@@ -2,7 +2,9 @@ package engine
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -44,6 +46,25 @@ func (o *outputWriter) Write(p []byte) (n int, err error) {
 	return len(p), nil
 }
 
+func compressEnv(envs []string) (result []string) {
+	for _, env := range envs {
+		k, v, ok := strings.Cut(env, "=")
+		if !ok || len(v) < 40_000 {
+			result = append(result, env)
+			continue
+		}
+
+		out := bytes.NewBuffer(nil)
+		b64 := base64.NewEncoder(base64.StdEncoding, out)
+		gz := gzip.NewWriter(b64)
+		_, _ = gz.Write([]byte(v))
+		_ = gz.Close()
+		_ = b64.Close()
+		result = append(result, k+`={"_gz":"`+out.String()+`"}`)
+	}
+	return
+}
+
 func (e *Engine) runCommand(ctx Context, tool types.Tool, input string, toolCategory ToolCategory) (cmdOut string, cmdErr error) {
 	id := counter.Next()
 
@@ -95,10 +116,10 @@ func (e *Engine) runCommand(ctx Context, tool types.Tool, input string, toolCate
 	for _, inputContext := range ctx.InputContext {
 		instructions = append(instructions, inputContext.Content)
 	}
+
 	var extraEnv = []string{
 		strings.TrimSpace("GPTSCRIPT_CONTEXT=" + strings.Join(instructions, "\n")),
 	}
-
 	cmd, stop, err := e.newCommand(ctx.Ctx, extraEnv, tool, input)
 	if err != nil {
 		return "", err
@@ -277,6 +298,6 @@ func (e *Engine) newCommand(ctx context.Context, extraEnv []string, tool types.T
 	}
 
 	cmd := exec.CommandContext(ctx, env.Lookup(envvars, args[0]), cmdArgs...)
-	cmd.Env = envvars
+	cmd.Env = compressEnv(envvars)
 	return cmd, stop, nil
 }

pkg/tests/runner_test.go

@@ -0,0 +1 @@
+{"_gz":"H4sIAAAAAAAA/+zAgQAAAADCMNb8JQK4wjYAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgHgAA//+94pKFQBkBAA=="}

pkg/tests/testdata/TestEnvOverflow/context.json

@@ -0,0 +1,14 @@
+context: c
+
+#!/bin/bash
+
+echo "${GPTSCRIPT_CONTEXT}"
+echo "${GPTSCRIPT_CONTEXT}" > ${GPTSCRIPT_TOOL_DIR}/context.json
+
+---
+name: c
+
+#!/bin/bash
+
+string=$(printf 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa%.0s' {1..1000})
+echo "$string"