Skip to content

Commit

Permalink
enhance: credentials: add GPTSCRIPT_CREDENTIAL_EXPIRATION (#709)
Browse files Browse the repository at this point in the history
Signed-off-by: Grant Linville <[email protected]>
  • Loading branch information
g-linville authored Aug 6, 2024
1 parent 160a733 commit b77cd13
Show file tree
Hide file tree
Showing 7 changed files with 97 additions and 3 deletions.
18 changes: 18 additions & 0 deletions docs/docs/03-tools/04-credential-tools.md
Original file line number Diff line number Diff line change
Expand Up @@ -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", ""))
```
24 changes: 21 additions & 3 deletions integration/cred_test.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
package integration

import (
"strings"
"testing"
"time"

"github.com/stretchr/testify/require"
)
Expand All @@ -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)
}
}
}
46 changes: 46 additions & 0 deletions integration/scripts/cred_expiration.gpt
Original file line number Diff line number Diff line change
@@ -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)
File renamed without changes.
2 changes: 2 additions & 0 deletions pkg/config/cliconfig.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions pkg/credentials/credential.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ const (
CredentialTypeTool CredentialType = "tool"
CredentialTypeModelProvider CredentialType = "modelProvider"
ExistingCredential = "GPTSCRIPT_EXISTING_CREDENTIAL"
CredentialExpiration = "GPTSCRIPT_CREDENTIAL_EXPIRATION"
)

type Credential struct {
Expand Down
9 changes: 9 additions & 0 deletions pkg/runner/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit b77cd13

Please sign in to comment.