feat(ateapi): a default egress policy on the ActorTemplate

cmd/ateapi/internal/controlapi/actor.go

@@ -79,6 +79,22 @@ func (s *RPCService) CreateActor(ctx context.Context, req *ateapipb.CreateActorR
 }
 
 func (s *ServiceImpl) CreateActor(ctx context.Context, inActor *ateapipb.Actor) (*ateapipb.Actor, error) {
+	return s.createActor(ctx, inActor, nil)
+}
+
+// CreateActorWithEgressPolicy creates the actor with the given "default"
+// egress policy instead of the one its template carries.
+func (s *ServiceImpl) CreateActorWithEgressPolicy(ctx context.Context, inActor *ateapipb.Actor, policy *ateapipb.EgressPolicy) (*ateapipb.Actor, error) {
+	if policy == nil {
+		return nil, status.Error(codes.InvalidArgument, "egress policy is required")
+	}
+	return s.createActor(ctx, inActor, policy)
+}
+
+// createActor records the actor and, when the template carries a
+// default_egress_policy or the caller passes one, the actor's "default"
+// egress policy in the same write. A passed policy wins over the template's.
+func (s *ServiceImpl) createActor(ctx context.Context, inActor *ateapipb.Actor, policy *ateapipb.EgressPolicy) (*ateapipb.Actor, error) {
 	// Check that the referenced ActorTemplate exists.
 	// FIXME: This is not atomic and it is not a guarantee that the template
 	// will still exist later.  Checking it here produces a nice error UX, but
@@ -149,8 +165,22 @@ func (s *ServiceImpl) CreateActor(ctx context.Context, inActor *ateapipb.Actor)
 		return nil, toGRPCInternalError(errs)
 	}
 
+	if policy == nil {
+		policy = egressPolicyFromTemplate(template, atespace)
+	}
+	if policy != nil {
+		if errs := Validate_EgressPolicy(ctx, operation.Operation{Type: operation.Create}, field.NewPath("egress_policy"), policy, nil); len(errs) > 0 {
+			return nil, toGRPCStatusError(errs)
+		}
+	}
+
 	// Save the data in the storage layer.
-	stored, err := s.store.CreateActor(ctx, outActor)
+	var stored *ateapipb.Actor
+	if policy != nil {
+		stored, err = s.store.CreateActorWithEgressPolicy(ctx, outActor, policy)
+	} else {
+		stored, err = s.store.CreateActor(ctx, outActor)
+	}
 	if err != nil {
 		if errors.Is(err, store.ErrAlreadyExists) {
 			return nil, status.Errorf(codes.AlreadyExists, "Actor %s already exists", name)
@@ -168,6 +198,19 @@ func (s *ServiceImpl) CreateActor(ctx context.Context, inActor *ateapipb.Actor)
 	return stored, nil
 }
 
+// egressPolicyFromTemplate returns the "default" EgressPolicy an actor in
+// atespace gets from its template, or nil when the template carries none.
+func egressPolicyFromTemplate(template *ateapipb.ActorTemplate, atespace string) *ateapipb.EgressPolicy {
+	defaultPolicy := template.GetDefaultEgressPolicy()
+	if defaultPolicy == nil {
+		return nil
+	}
+	return &ateapipb.EgressPolicy{
+		Metadata: &ateapipb.ResourceMetadata{Atespace: atespace, Name: "default"},
+		Rules:    proto.CloneOf(defaultPolicy).GetRules(),
+	}
+}
+
 // resolveTagSource resolves a CreateActor request's source tag
 // and checks that the tag is usable for creating an Actor in actorAtespace
 // from template.

cmd/ateapi/internal/controlapi/actor_test.go

@@ -16,6 +16,7 @@ package controlapi
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strings"
 	"testing"
@@ -1555,3 +1556,82 @@ func TestCreateActor_GoldenTagDefault(t *testing.T) {
 		})
 	}
 }
+
+// TestCreateActor_DefaultEgressPolicy pins that an actor created from a
+// template carrying default_egress_policy has its "default" EgressPolicy the
+// moment it exists, with the template's rules, that a policy the caller
+// passes wins over the template's, and that a template without one leaves
+// the actor without a policy.
+func TestCreateActor_DefaultEgressPolicy(t *testing.T) {
+	rules := []*ateapipb.EgressRule{
+		{Hostnames: &ateapipb.HostnameRule{Patterns: []string{"api.anthropic.com"}}},
+		{Hostnames: &ateapipb.HostnameRule{Patterns: []string{"github.com"}}},
+	}
+	for _, scenario := range []string{"with policy", "without policy", "explicit policy wins"} {
+		t.Run(scenario, func(t *testing.T) {
+			ctx := t.Context()
+			persistence := newTestPersistence(t)
+			storetest.MustCreateAtespace(t, ctx, persistence, "team-a")
+			tmpl := seedSubstrateTemplate(t, ctx, persistence, "tmpl")
+			if scenario != "without policy" {
+				var err error
+				tmpl, err = persistence.UpdateActorTemplate(ctx, resources.ActorTemplateRefFromActorTemplate(tmpl), store.PreconditionFrom(tmpl), func(db *ateapipb.ActorTemplate) error {
+					db.DefaultEgressPolicy = &ateapipb.EgressPolicyTemplate{Rules: rules}
+					return nil
+				})
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+			actor := &ateapipb.Actor{Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "actor"}, ActorTemplate: resources.ActorTemplateRefFromActorTemplate(tmpl).ToObjectRef()}
+			svc := &ServiceImpl{store: persistence}
+			var created *ateapipb.Actor
+			var err error
+			wantRules := rules
+			switch scenario {
+			case "explicit policy wins":
+				wantRules = []*ateapipb.EgressRule{{Cidrs: &ateapipb.CIDRRule{Cidrs: []string{"10.0.0.0/8"}}}}
+				created, err = svc.CreateActorWithEgressPolicy(ctx, actor, &ateapipb.EgressPolicy{
+					Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "default"},
+					Rules:    wantRules,
+				})
+			default:
+				created, err = svc.CreateActor(ctx, actor)
+			}
+			if err != nil {
+				t.Fatalf("CreateActor = %v", err)
+			}
+			policy, err := persistence.GetEgressPolicy(ctx, resources.ActorRefFromActor(created))
+			if scenario == "without policy" {
+				if !errors.Is(err, store.ErrNotFound) {
+					t.Fatalf("GetEgressPolicy = %v, %v; want ErrNotFound", policy, err)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("GetEgressPolicy = %v", err)
+			}
+			if got := policy.GetMetadata(); got.GetName() != "default" || got.GetAtespace() != "team-a" {
+				t.Fatalf("policy metadata = %v", got)
+			}
+			if diff := cmp.Diff(wantRules, policy.GetRules(), protocmp.Transform()); diff != "" {
+				t.Fatalf("policy rules differ (-want +got):\n%s", diff)
+			}
+			// The copy is a snapshot: a later change to the template's default
+			// does not reach the actor's policy.
+			if _, err := persistence.UpdateActorTemplate(ctx, resources.ActorTemplateRefFromActorTemplate(tmpl), store.PreconditionFrom(tmpl), func(db *ateapipb.ActorTemplate) error {
+				db.DefaultEgressPolicy = &ateapipb.EgressPolicyTemplate{Rules: []*ateapipb.EgressRule{{Hostnames: &ateapipb.HostnameRule{Patterns: []string{"changed.example.com"}}}}}
+				return nil
+			}); err != nil {
+				t.Fatal(err)
+			}
+			again, err := persistence.GetEgressPolicy(ctx, resources.ActorRefFromActor(created))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if diff := cmp.Diff(wantRules, again.GetRules(), protocmp.Transform()); diff != "" {
+				t.Fatalf("stored rules changed with the template (-want +got):\n%s", diff)
+			}
+		})
+	}
+}

cmd/ateapi/internal/store/atepg/actor.go

@@ -27,6 +27,34 @@ import (
 )
 
 func (p *Persistence) CreateActor(ctx context.Context, actor *ateapipb.Actor) (*ateapipb.Actor, error) {
+	return insertActor(ctx, p.pool, actor)
+}
+
+// CreateActorWithEgressPolicy inserts the actor and its "default" egress
+// policy in one transaction, so no observer ever sees the actor without the
+// policy. The policy's foreign key on the actor row is satisfied within the
+// transaction.
+func (p *Persistence) CreateActorWithEgressPolicy(ctx context.Context, actor *ateapipb.Actor, policy *ateapipb.EgressPolicy) (*ateapipb.Actor, error) {
+	tx, err := p.pool.Begin(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("beginning transaction: %w", err)
+	}
+	defer tx.Rollback(ctx) //nolint:errcheck // no-op once committed
+
+	dbActor, err := insertActor(ctx, tx, actor)
+	if err != nil {
+		return nil, err
+	}
+	if _, err := insertEgressPolicy(ctx, tx, resources.ActorRefFromActor(dbActor), policy); err != nil {
+		return nil, err
+	}
+	if err := tx.Commit(ctx); err != nil {
+		return nil, fmt.Errorf("committing actor with egress policy: %w", err)
+	}
+	return dbActor, nil
+}
+
+func insertActor(ctx context.Context, q querier, actor *ateapipb.Actor) (*ateapipb.Actor, error) {
 	atespace := actor.GetMetadata().GetAtespace()
 	name := actor.GetMetadata().GetName()
 
@@ -42,7 +70,7 @@ func (p *Persistence) CreateActor(ctx context.Context, actor *ateapipb.Actor) (*
 		return nil, fmt.Errorf("marshaling actor: %w", err)
 	}
 
-	_, err = p.pool.Exec(ctx, `
+	_, err = q.Exec(ctx, `
 		INSERT INTO actors (atespace, name, uid, version, proto)
 		VALUES ($1, $2, $3, $4, $5)`,
 		atespace, name, dbActor.GetMetadata().GetUid(), dbActor.GetMetadata().GetVersion(), protoBytes)

cmd/ateapi/internal/store/atepg/egress_policy.go

@@ -27,13 +27,17 @@ import (
 )
 
 func (p *Persistence) CreateEgressPolicy(ctx context.Context, actorRef resources.ActorRef, policy *ateapipb.EgressPolicy) (*ateapipb.EgressPolicy, error) {
+	return insertEgressPolicy(ctx, p.pool, actorRef, policy)
+}
+
+func insertEgressPolicy(ctx context.Context, q querier, actorRef resources.ActorRef, policy *ateapipb.EgressPolicy) (*ateapipb.EgressPolicy, error) {
 	dbPolicy := proto.Clone(policy).(*ateapipb.EgressPolicy)
 	dbPolicy.Metadata = newCreateMetadata(actorRef.Atespace, "default")
 	protoBytes, err := proto.Marshal(dbPolicy)
 	if err != nil {
 		return nil, fmt.Errorf("marshaling egress policy: %w", err)
 	}
-	_, err = p.pool.Exec(ctx, `
+	_, err = q.Exec(ctx, `
 		INSERT INTO actor_egress_policies (atespace, actor_name, uid, version, proto)
 		VALUES ($1, $2, $3, $4, $5)`, actorRef.Atespace, actorRef.Name, dbPolicy.GetMetadata().GetUid(), dbPolicy.GetMetadata().GetVersion(), protoBytes)
 	if err != nil {

cmd/ateapi/internal/store/store.go

@@ -72,6 +72,11 @@ type Interface interface {
 	// ErrFailedPrecondition if the actor's atespace does not exist.
 	CreateActor(ctx context.Context, actor *ateapipb.Actor) (*ateapipb.Actor, error)
 
+	// CreateActorWithEgressPolicy is CreateActor plus the actor's "default"
+	// egress policy, written in the same transaction: either both exist
+	// afterwards or neither does. Returns the same errors as CreateActor.
+	CreateActorWithEgressPolicy(ctx context.Context, actor *ateapipb.Actor, policy *ateapipb.EgressPolicy) (*ateapipb.Actor, error)
+
 	// Fetches an actor by reference. Returns ErrNotFound if missing.
 	GetActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error)
 

cmd/ateapi/internal/store/storecontract/contract.go

@@ -218,6 +218,56 @@ func RunContractTests(t *testing.T, setup func(t *testing.T) store.Interface) {
 func runEgressPolicyContractTests(t *testing.T, setup func(t *testing.T) store.Interface) {
 	t.Helper()
 
+	t.Run("CreateActorWithEgressPolicy", func(t *testing.T) {
+		s := setup(t)
+		ctx := context.Background()
+		mustCreateAtespace(t, s, testAtespace)
+		actor := &ateapipb.Actor{
+			Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "session-1"},
+			Status:   &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_SUSPENDED},
+		}
+		policy := &ateapipb.EgressPolicy{Rules: []*ateapipb.EgressRule{{
+			Hostnames: &ateapipb.HostnameRule{Patterns: []string{"api.example.com"}},
+		}}}
+
+		created, err := s.CreateActorWithEgressPolicy(ctx, actor, policy)
+		if err != nil {
+			t.Fatalf("CreateActorWithEgressPolicy failed: %v", err)
+		}
+		actorRef := resources.ActorRefFromActor(created)
+		if created.GetMetadata().GetUid() == "" || created.GetMetadata().GetVersion() != 1 {
+			t.Fatalf("created actor metadata = %v", created.GetMetadata())
+		}
+		got, err := s.GetEgressPolicy(ctx, actorRef)
+		if err != nil {
+			t.Fatalf("GetEgressPolicy failed: %v", err)
+		}
+		if md := got.GetMetadata(); md.GetName() != "default" || md.GetAtespace() != testAtespace || md.GetUid() == "" || md.GetVersion() != 1 {
+			t.Fatalf("policy metadata = %v", md)
+		}
+		if diff := cmp.Diff(policy.GetRules(), got.GetRules(), protocmp.Transform()); diff != "" {
+			t.Fatalf("policy rules differ (-want +got):\n%s", diff)
+		}
+		if _, err := s.CreateActorWithEgressPolicy(ctx, actor, policy); !errors.Is(err, store.ErrAlreadyExists) {
+			t.Fatalf("duplicate create error = %v, want ErrAlreadyExists", err)
+		}
+
+		// A failed write leaves neither row behind.
+		missing := &ateapipb.Actor{
+			Metadata: &ateapipb.ResourceMetadata{Atespace: "no-such-atespace", Name: "session-2"},
+			Status:   &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_SUSPENDED},
+		}
+		if _, err := s.CreateActorWithEgressPolicy(ctx, missing, policy); !errors.Is(err, store.ErrFailedPrecondition) {
+			t.Fatalf("create in missing atespace error = %v, want ErrFailedPrecondition", err)
+		}
+		if _, err := s.GetActor(ctx, resources.ActorRefFromActor(missing)); !errors.Is(err, store.ErrNotFound) {
+			t.Fatalf("GetActor after failed create = %v, want ErrNotFound", err)
+		}
+		if _, err := s.GetEgressPolicy(ctx, resources.ActorRefFromActor(missing)); !errors.Is(err, store.ErrNotFound) {
+			t.Fatalf("GetEgressPolicy after failed create = %v, want ErrNotFound", err)
+		}
+	})
+
 	t.Run("EgressPolicy_Lifecycle", func(t *testing.T) {
 		s := setup(t)
 		ctx := context.Background()