enhance: credentials: add GPTSCRIPT_CREDENTIAL_EXPIRATION (#709)

Signed-off-by: Grant Linville <[email protected]>

Author: g-linville

docs/docs/03-tools/04-credential-tools.md

@@ -204,3 +204,21 @@ that environment variable, and if it is set, get the refresh token from the exis
 typically without user interaction.
 
 For an example of a tool that uses the refresh feature, see the [Gateway OAuth2 tool](https://github.com/gptscript-ai/gateway-oauth2).
+
+### GPTSCRIPT_CREDENTIAL_EXPIRATION environment variable
+
+When a tool references a credential tool, GPTScript will add the environment variables from the credential to the tool's
+environment before executing the tool. If at least one of the credentials has an `expiresAt` field, GPTScript will also
+set the environment variable `GPTSCRIPT_CREDENTIAL_EXPIRATION`, which contains the nearest expiration time out of all
+credentials referenced by the tool, in RFC 3339 format. That way, it can be referenced in the tool body if needed.
+Here is an example:
+
+```
+Credential: my-credential-tool.gpt as myCred
+
+#!python3
+
+import os
+
+print("myCred expires at " + os.getenv("GPTSCRIPT_CREDENTIAL_EXPIRATION", ""))
+```

integration/cred_test.go

@@ -1,7 +1,9 @@
 package integration
 
 import (
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
@@ -15,15 +17,31 @@ func TestGPTScriptCredential(t *testing.T) {
 // TestCredentialScopes makes sure that environment variables set by credential tools and shared credential tools
 // are only available to the correct tools. See scripts/credscopes.gpt for more details.
 func TestCredentialScopes(t *testing.T) {
-	out, err := RunScript("scripts/credscopes.gpt", "--sub-tool", "oneOne")
+	out, err := RunScript("scripts/cred_scopes.gpt", "--sub-tool", "oneOne")
 	require.NoError(t, err)
 	require.Contains(t, out, "good")
 
-	out, err = RunScript("scripts/credscopes.gpt", "--sub-tool", "twoOne")
+	out, err = RunScript("scripts/cred_scopes.gpt", "--sub-tool", "twoOne")
 	require.NoError(t, err)
 	require.Contains(t, out, "good")
 
-	out, err = RunScript("scripts/credscopes.gpt", "--sub-tool", "twoTwo")
+	out, err = RunScript("scripts/cred_scopes.gpt", "--sub-tool", "twoTwo")
 	require.NoError(t, err)
 	require.Contains(t, out, "good")
 }
+
+// TestCredentialExpirationEnv tests a GPTScript with two credentials that expire at different times.
+// One expires after two hours, and the other expires after one hour.
+// This test makes sure that the GPTSCRIPT_CREDENTIAL_EXPIRATION environment variable is set to the nearer expiration time (1h).
+func TestCredentialExpirationEnv(t *testing.T) {
+	out, err := RunScript("scripts/cred_expiration.gpt")
+	require.NoError(t, err)
+
+	for _, line := range strings.Split(out, "\n") {
+		if timestamp, found := strings.CutPrefix(line, "Expires: "); found {
+			expiresTime, err := time.Parse(time.RFC3339, timestamp)
+			require.NoError(t, err)
+			require.True(t, time.Until(expiresTime) < time.Hour)
+		}
+	}
+}

integration/scripts/cred_expiration.gpt

@@ -0,0 +1,46 @@
+cred: credentialTool with 2 as hours
+cred: credentialTool with 1 as hours
+
+#!python3
+
+import os
+
+print("Expires: " + os.getenv("GPTSCRIPT_CREDENTIAL_EXPIRATION", ""), end="")
+
+---
+name: credentialTool
+args: hours: the number of hours from now to expire
+
+#!python3
+
+import os
+import json
+from datetime import datetime, timedelta, timezone
+
+class Output:
+    def __init__(self, env, expires_at):
+        self.env = env
+        self.expiresAt = expires_at
+
+    def to_dict(self):
+        return {
+            "env": self.env,
+            "expiresAt": self.expiresAt.isoformat()
+        }
+
+hours_str = os.getenv("HOURS")
+if hours_str is None:
+	print("HOURS environment variable is not set")
+	os._exit(1)
+
+try:
+	hours = int(hours_str)
+except ValueError:
+	print("failed to parse HOURS")
+	os._exit(1)
+
+expires_at = datetime.now(timezone.utc) + timedelta(hours=hours)
+out = Output(env={"yeet": "yote"}, expires_at=expires_at)
+out_json = json.dumps(out.to_dict())
+
+print(out_json)

integration/scripts/credscopes.gpt renamed to integration/scripts/cred_scopes.gpt

@@ -55,6 +55,8 @@ type CLIConfig struct {
 	Auths               map[string]AuthConfig `json:"auths,omitempty"`
 	CredentialsStore    string                `json:"credsStore,omitempty"`
 	GPTScriptConfigFile string                `json:"gptscriptConfig,omitempty"`
+	GatewayURL          string                `json:"gatewayURL,omitempty"`
+	Integrations        map[string]string     `json:"integrations,omitempty"`
 
 	auths     map[string]types.AuthConfig
 	authsLock *sync.Mutex

pkg/config/cliconfig.go

@@ -16,6 +16,7 @@ const (
 	CredentialTypeTool          CredentialType = "tool"
 	CredentialTypeModelProvider CredentialType = "modelProvider"
 	ExistingCredential                         = "GPTSCRIPT_EXISTING_CREDENTIAL"
+	CredentialExpiration                       = "GPTSCRIPT_CREDENTIAL_EXPIRATION"
 )
 
 type Credential struct {

pkg/credentials/credential.go

@@ -865,6 +865,7 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
 		}
 	}
 
+	var nearestExpiration *time.Time
 	for _, ref := range credToolRefs {
 		toolName, credentialAlias, args, err := types.ParseCredentialArgs(ref.Reference, callCtx.Input)
 		if err != nil {
@@ -967,11 +968,19 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
 			} else {
 				log.Warnf("Not saving credential for tool %s - credentials will only be saved for tools from GitHub, or tools that use aliases.", toolName)
 			}
+
+			if c.ExpiresAt != nil && (nearestExpiration == nil || nearestExpiration.After(*c.ExpiresAt)) {
+				nearestExpiration = c.ExpiresAt
+			}
 		}
 
 		for k, v := range c.Env {
 			env = append(env, fmt.Sprintf("%s=%s", k, v))
 		}
+
+		if nearestExpiration != nil {
+			env = append(env, fmt.Sprintf("%s=%s", credentials.CredentialExpiration, nearestExpiration.Format(time.RFC3339)))
+		}
 	}
 
 	return env, nil