Scope an init script to the environment it targets (#103)

Author: vitramir

internal/server/environment_authorization_db_test.go

@@ -50,6 +50,9 @@ func seedEnvironment(ctx context.Context, t *testing.T, backing *store.Store, or
 		RunnerID:     &runnerID,
 		Flavor:       "small",
 		Availability: store.EnvironmentAvailabilityPrivate,
+		LLMMode:      store.LLMModePlatform,
+		// NOT NULL, and a nil slice encodes as NULL rather than '{}'.
+		LLMAllowedModels: []string{},
 	})
 	if err != nil {
 		t.Fatalf("seed environment: %v", err)

internal/server/init_script_environment_db_test.go

@@ -0,0 +1,86 @@
+package server
+
+import (
+	"context"
+	"testing"
+
+	"github.com/agynio/agents/internal/store"
+	"github.com/google/uuid"
+)
+
+// An init script targets an agent, an mcp or an environment. Only the first two
+// were ever filtered on, so a list scoped to one environment returned every
+// script in the database -- and every mutation tried to resolve an agent a
+// environment-scoped row does not have, which rolled the write back.
+func seedEnvironmentInitScript(ctx context.Context, t *testing.T, backing *store.Store, organizationID, environmentID uuid.UUID, script string) uuid.UUID {
+	t.Helper()
+	created, err := backing.CreateInitScript(ctx, organizationID, store.InitScriptInput{
+		Script:        script,
+		Description:   script,
+		EnvironmentID: &environmentID,
+	})
+	if err != nil {
+		t.Fatalf("create init script: %v", err)
+	}
+	return created.Meta.ID
+}
+
+func TestListInitScriptsIsScopedToItsEnvironment(t *testing.T) {
+	ctx := context.Background()
+	_, backing := environmentAuthorizationServer(ctx, t, &recordingAuthorizationWriter{})
+	organizationID := uuid.New()
+	first := seedEnvironment(ctx, t, backing, organizationID, "first")
+	second := seedEnvironment(ctx, t, backing, organizationID, "second")
+	seedEnvironmentInitScript(ctx, t, backing, organizationID, first, "echo first")
+
+	result, err := backing.ListInitScripts(ctx, store.InitScriptFilter{EnvironmentID: &first}, 100, nil)
+	if err != nil {
+		t.Fatalf("list first: %v", err)
+	}
+	if len(result.InitScripts) != 1 {
+		t.Fatalf("expected the environment's own script, got %d", len(result.InitScripts))
+	}
+
+	// The leak: an unfiltered query returns this one too.
+	other, err := backing.ListInitScripts(ctx, store.InitScriptFilter{EnvironmentID: &second}, 100, nil)
+	if err != nil {
+		t.Fatalf("list second: %v", err)
+	}
+	if len(other.InitScripts) != 0 {
+		t.Fatalf("expected no scripts for an environment that has none, got %d", len(other.InitScripts))
+	}
+}
+
+func TestDeleteEnvironmentInitScript(t *testing.T) {
+	ctx := context.Background()
+	_, backing := environmentAuthorizationServer(ctx, t, &recordingAuthorizationWriter{})
+	organizationID := uuid.New()
+	environmentID := seedEnvironment(ctx, t, backing, organizationID, "first")
+	scriptID := seedEnvironmentInitScript(ctx, t, backing, organizationID, environmentID, "echo first")
+
+	// Resolving an agent for a row that has none used to fail inside the
+	// transaction, so the delete rolled back and reported an internal error.
+	if err := backing.DeleteInitScript(ctx, scriptID); err != nil {
+		t.Fatalf("delete init script: %v", err)
+	}
+	result, err := backing.ListInitScripts(ctx, store.InitScriptFilter{EnvironmentID: &environmentID}, 100, nil)
+	if err != nil {
+		t.Fatalf("list after delete: %v", err)
+	}
+	if len(result.InitScripts) != 0 {
+		t.Fatalf("expected the script to be gone, got %d", len(result.InitScripts))
+	}
+}
+
+func TestUpdateEnvironmentInitScript(t *testing.T) {
+	ctx := context.Background()
+	_, backing := environmentAuthorizationServer(ctx, t, &recordingAuthorizationWriter{})
+	organizationID := uuid.New()
+	environmentID := seedEnvironment(ctx, t, backing, organizationID, "first")
+	scriptID := seedEnvironmentInitScript(ctx, t, backing, organizationID, environmentID, "echo first")
+
+	updated := "echo second"
+	if _, err := backing.UpdateInitScript(ctx, scriptID, store.InitScriptUpdate{Script: &updated}); err != nil {
+		t.Fatalf("update init script: %v", err)
+	}
+}

internal/store/store.go

@@ -1127,11 +1127,7 @@ func (s *Store) CreateInitScript(ctx context.Context, organizationID uuid.UUID,
 			}
 			return InitScript{}, err
 		}
-		agentID, err := resolveAgentID(ctx, tx, script.AgentID, script.McpID)
-		if err != nil {
-			return InitScript{}, err
-		}
-		if err := touchAgentUpdatedAt(ctx, tx, agentID); err != nil {
+		if err := touchTargetAgent(ctx, tx, script.AgentID, script.McpID, script.EnvironmentID); err != nil {
 			return InitScript{}, err
 		}
 		return script, nil
@@ -1175,11 +1171,7 @@ func (s *Store) UpdateInitScript(ctx context.Context, id uuid.UUID, update InitS
 			}
 			return InitScript{}, err
 		}
-		agentID, err := resolveAgentID(ctx, tx, script.AgentID, script.McpID)
-		if err != nil {
-			return InitScript{}, err
-		}
-		if err := touchAgentUpdatedAt(ctx, tx, agentID); err != nil {
+		if err := touchTargetAgent(ctx, tx, script.AgentID, script.McpID, script.EnvironmentID); err != nil {
 			return InitScript{}, err
 		}
 		return script, nil
@@ -1199,11 +1191,7 @@ func (s *Store) DeleteInitScript(ctx context.Context, id uuid.UUID) error {
 			}
 			return struct{}{}, err
 		}
-		agentID, err := resolveAgentID(ctx, tx, script.AgentID, script.McpID)
-		if err != nil {
-			return struct{}{}, err
-		}
-		if err := touchAgentUpdatedAt(ctx, tx, agentID); err != nil {
+		if err := touchTargetAgent(ctx, tx, script.AgentID, script.McpID, script.EnvironmentID); err != nil {
 			return struct{}{}, err
 		}
 		return struct{}{}, nil
@@ -1220,6 +1208,9 @@ func (s *Store) ListInitScripts(ctx context.Context, filter InitScriptFilter, pa
 	if filter.McpID != nil {
 		clauses, args = appendClause(clauses, args, "mcp_id = $%d", *filter.McpID)
 	}
+	if filter.EnvironmentID != nil {
+		clauses, args = appendClause(clauses, args, "environment_id = $%d", *filter.EnvironmentID)
+	}
 
 	scripts, nextCursor, err := listEntities(ctx, s.pool,
 		fmt.Sprintf("SELECT %s FROM init_scripts", initScriptColumns),