modify schema for environments in tool call proxy

move ci env

switch environment parameter for tool proxy

Author: qstearns

server/internal/environments/shared.go

@@ -66,6 +66,31 @@ func (e *EnvironmentEntries) Load(ctx context.Context, projectID uuid.UUID, envI
 	return envMap, nil
 }
 
+func (e *EnvironmentEntries) LoadSourceEnv(ctx context.Context, projectID uuid.UUID, sourceKind string, sourceSlug string) (map[string]string, error) {
+	sourceEnv, err := e.repo.GetEnvironmentForSource(ctx, repo.GetEnvironmentForSourceParams{
+		SourceKind: sourceKind,
+		SourceSlug: sourceSlug,
+		ProjectID:  projectID,
+	})
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return map[string]string{}, nil
+		}
+		return nil, fmt.Errorf("get environment for source: %w", err)
+	}
+
+	entries, err := e.ListEnvironmentEntries(ctx, projectID, sourceEnv.ID, false)
+	if err != nil {
+		return nil, fmt.Errorf("list environment entries: %w", err)
+	}
+
+	envMap := make(map[string]string, len(entries))
+	for _, entry := range entries {
+		envMap[entry.Name] = entry.Value
+	}
+	return envMap, nil
+}
+
 func (e *EnvironmentEntries) ListEnvironmentEntries(ctx context.Context, projectID uuid.UUID, environmentID uuid.UUID, redacted bool) ([]repo.EnvironmentEntry, error) {
 	entries, err := e.repo.ListEnvironmentEntries(ctx, repo.ListEnvironmentEntriesParams{
 		ProjectID:     projectID,

server/internal/gateway/env.go

@@ -3,6 +3,7 @@ package gateway
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/google/uuid"
 )
@@ -14,6 +15,13 @@ type EnvironmentLoader interface {
 	//   * [ErrNotFound]: when the environment does not exist.
 	//   * `error`: when an unrecognized error occurs.
 	Load(ctx context.Context, projectID uuid.UUID, environmentID SlugOrID) (map[string]string, error)
+
+	// LoadSourceEnv retrieves the environment variables associated with a tool's source.
+	// Returns an empty map if no source environment exists.
+	//
+	// # Errors
+	//   * `error`: when an unrecognized error occurs.
+	LoadSourceEnv(ctx context.Context, projectID uuid.UUID, sourceKind string, sourceSlug string) (map[string]string, error)
 }
 
 type SlugOrID struct {
@@ -51,3 +59,32 @@ func (s *SlugOrID) String() string {
 func (s *SlugOrID) IsEmpty() bool {
 	return s.ID == uuid.Nil && s.Slug == ""
 }
+
+type CaseInsensitiveEnv struct {
+	data map[string]string
+}
+
+func NewCaseInsensitiveEnv() *CaseInsensitiveEnv {
+	return &CaseInsensitiveEnv{data: make(map[string]string)}
+}
+
+func CIEnvFrom(vars map[string]string) *CaseInsensitiveEnv {
+	env := NewCaseInsensitiveEnv()
+	for k, v := range vars {
+		env.Set(k, v)
+	}
+	return env
+}
+
+func (c *CaseInsensitiveEnv) Get(key string) string {
+	return c.data[strings.ToLower(key)]
+}
+
+func (c *CaseInsensitiveEnv) Set(key, value string) {
+	c.data[strings.ToLower(key)] = value
+}
+
+type ToolCallEnv struct {
+	SystemEnv  *CaseInsensitiveEnv
+	UserConfig *CaseInsensitiveEnv
+}

server/internal/gateway/proxy.go

@@ -74,26 +74,6 @@ type ToolCallBody struct {
 	GramRequestSummary   string            `json:"gram-request-summary"`
 }
 
-// CaseInsensitiveEnv provides case-insensitive environment variable lookup.
-type CaseInsensitiveEnv struct {
-	data map[string]string
-}
-
-// NewCaseInsensitiveEnv creates a new empty case-insensitive environment.
-func NewCaseInsensitiveEnv() *CaseInsensitiveEnv {
-	return &CaseInsensitiveEnv{data: make(map[string]string)}
-}
-
-// Get retrieves an environment variable value by key (case-insensitive).
-func (c *CaseInsensitiveEnv) Get(key string) string {
-	return c.data[strings.ToLower(key)]
-}
-
-// Set sets an environment variable value by key (case-insensitive).
-func (c *CaseInsensitiveEnv) Set(key, value string) {
-	c.data[strings.ToLower(key)] = value
-}
-
 type toolcallErrorSchema struct {
 	Error string `json:"error"`
 }
@@ -146,7 +126,7 @@ func (tp *ToolProxy) Do(
 	ctx context.Context,
 	w http.ResponseWriter,
 	requestBody io.Reader,
-	env *CaseInsensitiveEnv,
+	env ToolCallEnv,
 	plan *ToolCallPlan,
 	toolCallLogger tm.ToolCallLogger,
 ) (err error) {
@@ -172,10 +152,6 @@ func (tp *ToolProxy) Do(
 		attr.SlogToolCallSource(string(tp.source)),
 	)
 
-	if env == nil {
-		env = NewCaseInsensitiveEnv()
-	}
-
 	switch plan.Kind {
 	case "":
 		return oops.E(oops.CodeInvariantViolation, nil, "tool kind is not set").Log(ctx, tp.logger)
@@ -195,7 +171,7 @@ func (tp *ToolProxy) doFunction(
 	logger *slog.Logger,
 	w http.ResponseWriter,
 	requestBody io.Reader,
-	env *CaseInsensitiveEnv,
+	env ToolCallEnv,
 	descriptor *ToolDescriptor,
 	plan *FunctionToolCallPlan,
 	toolCallLogger tm.ToolCallLogger,
@@ -307,7 +283,7 @@ func (tp *ToolProxy) doHTTP(
 	logger *slog.Logger,
 	w http.ResponseWriter,
 	requestBody io.Reader,
-	env *CaseInsensitiveEnv,
+	env ToolCallEnv,
 	descriptor *ToolDescriptor,
 	plan *HTTPToolCallPlan,
 	toolCallLogger tm.ToolCallLogger,
@@ -570,7 +546,7 @@ type promptGetParams struct {
 	Arguments map[string]any `json:"arguments"`
 }
 
-func (tp *ToolProxy) doPrompt(ctx context.Context, logger *slog.Logger, w http.ResponseWriter, requestBody io.Reader, env *CaseInsensitiveEnv, descriptor *ToolDescriptor, plan *PromptToolCallPlan) error {
+func (tp *ToolProxy) doPrompt(ctx context.Context, logger *slog.Logger, w http.ResponseWriter, requestBody io.Reader, env ToolCallEnv, descriptor *ToolDescriptor, plan *PromptToolCallPlan) error {
 	var params promptGetParams
 	if err := json.NewDecoder(requestBody).Decode(&params); err != nil {
 		return oops.E(oops.CodeBadRequest, err, "failed to parse get prompt request").Log(ctx, logger)
@@ -622,7 +598,7 @@ func retryWithBackoff(
 		if err != nil {
 			continue
 		}
-
+		// check if we should retry based on method and status code
 		if !slices.Contains(retryBackoff.methods, resp.Request.Method) || !slices.Contains(retryBackoff.statusCodes, resp.StatusCode) {
 			return resp, err
 		}

server/internal/gateway/security.go

@@ -29,7 +29,7 @@ func processSecurity(
 	tool *ToolDescriptor,
 	plan *HTTPToolCallPlan,
 	cacheImpl cache.Cache,
-	envVars *CaseInsensitiveEnv,
+	env ToolCallEnv,
 	serverURL string,
 	toolCallLogger tm.ToolCallLogger,
 ) bool {

server/internal/mcp/rpc_tools_call.go

@@ -361,7 +361,8 @@ func resolveUserConfiguration(
 ) (map[string]string, error) {
 	userConfig := make(map[string]string)
 
-	// IMPORTANT: MCP servers accessed in a public manner or not gram authenticated, there is no concept of using stored environments for them
+	// IMPORTANT: we must only attach gram environments to authenticated payloads. Gram environments contain
+	// secrets owned by Gram projects and should not be usable by public clients
 	if payload.environment != "" && payload.authenticated {
 		storedEnvVars, err := env.Load(ctx, payload.projectID, gateway.Slug(payload.environment))
 		switch {