agentplane: make sampler RPC handlers owner-private

agentplane/central_sampler/BUILD.bazel

@@ -326,10 +326,8 @@ go_library(
         "executor_session.go",
         "protocol_adapter.go",
         "provider_recovery_runtime.go",
-        "provider_request_read.go",
         "runtime.go",
         "runtime_replica.go",
-        "service.go",
         "session_compaction.go",
         "session_driver.go",
         "session_host.go",
@@ -341,7 +339,10 @@ go_library(
         "view_image.go",
     ],
     importpath = "github.com/agentydragon/ducktape/agentplane/central_sampler",
-    visibility = ["//agentplane/central_sampler/cmd:__pkg__"],
+    visibility = [
+        "//agentplane/central_sampler/cmd:__pkg__",
+        "//agentplane/central_sampler/internal/rpc:__pkg__",
+    ],
     deps = [
         "//agentplane/central_sampler/api:sampler_go_proto",
         "//agentplane/central_sampler/historyexport",
@@ -400,9 +401,7 @@ go_test(
         "provider_finish_runtime_test.go",
         "provider_recovery_runtime_test.go",
         "provider_turn_progress_test.go",
-        "request_read_service_test.go",
         "runtime_lifecycle_test.go",
-        "service_test.go",
         "settings_fixture_test.go",
         "unified_exec_test.go",
     ],

agentplane/central_sampler/cmd/BUILD.bazel

@@ -15,6 +15,7 @@ go_library(
         "//agentplane/central_sampler",
         "//agentplane/central_sampler/api:sampler_go_proto",
         "//agentplane/central_sampler/historyexport",
+        "//agentplane/central_sampler/internal/rpc",
         "//agentplane/central_sampler/internal/sessionstore",
         "//agentplane/central_sampler/internal/settings",
         "//agentplane/central_sampler/migrations",

agentplane/central_sampler/cmd/main.go

@@ -15,6 +15,7 @@ import (
 	centralsampler "github.com/agentydragon/ducktape/agentplane/central_sampler"
 	samplerpb "github.com/agentydragon/ducktape/agentplane/central_sampler/api/samplerpb"
 	"github.com/agentydragon/ducktape/agentplane/central_sampler/historyexport"
+	samplerrpc "github.com/agentydragon/ducktape/agentplane/central_sampler/internal/rpc"
 	"github.com/agentydragon/ducktape/agentplane/central_sampler/internal/sessionstore"
 	samplersettings "github.com/agentydragon/ducktape/agentplane/central_sampler/internal/settings"
 	"github.com/agentydragon/ducktape/agentplane/central_sampler/migrations"
@@ -130,7 +131,7 @@ func run(ctx context.Context, settings samplersettings.Settings) error {
 	defer unsubscribe()
 	runtime := centralsampler.NewRuntime(store, settings, centralsampler.NewCodeModeHost(context.Background(), settings, store))
 	server := grpc.NewServer(grpc.MaxRecvMsgSize(64*1024*1024), grpc.MaxSendMsgSize(64*1024*1024))
-	samplerpb.RegisterSamplerServer(server, &centralsampler.Sampler{Store: store, Runtime: runtime, Resolver: resolver, Settings: settings, Wakeups: hub})
+	samplerpb.RegisterSamplerServer(server, &samplerrpc.Sampler{Store: store, Runtime: runtime, Resolver: resolver, Settings: settings, Wakeups: hub})
 	healthServer := health.NewServer()
 	healthpb.RegisterHealthServer(server, healthServer)
 	listener, err := netlistener.Open(settings.Listen)

agentplane/central_sampler/internal/rpc/BUILD.bazel

@@ -0,0 +1,87 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+package(default_visibility = ["//visibility:private"])
+
+go_library(
+    name = "rpc",
+    srcs = [
+        "provider_request_read.go",
+        "service.go",
+    ],
+    importpath = "github.com/agentydragon/ducktape/agentplane/central_sampler/internal/rpc",
+    visibility = ["//agentplane/central_sampler/cmd:__pkg__"],
+    deps = [
+        "//agentplane/central_sampler",
+        "//agentplane/central_sampler/api:sampler_go_proto",
+        "//agentplane/central_sampler/internal/sessionstore",
+        "//agentplane/central_sampler/internal/settings",
+        "//agentplane/central_sampler/protocol",
+        "//agentplane/central_sampler/providers/responses",
+        "//agentplane/env/runtime/control:environment_runtime_go_proto",
+        "//agentplane/identity/auth",
+        "//agentplane/postgres/wakeups",
+        "//agentplane/proto:command_go_proto",
+        "//agentplane/proto:event_log_go_proto",
+        "//agentplane/proto:session_go_proto",
+        "@com_github_google_uuid//:uuid",
+        "@org_golang_google_grpc//:grpc",
+        "@org_golang_google_grpc//codes",
+        "@org_golang_google_grpc//metadata",
+        "@org_golang_google_grpc//status",
+    ],
+)
+
+go_test(
+    name = "unit_test",
+    size = "medium",
+    srcs = [
+        "service_test.go",
+        "test_helpers_test.go",
+    ],
+    embed = [":rpc"],
+    race = "on",
+    deps = [
+        "//agentplane/central_sampler/internal/sessionstore",
+        "//agentplane/central_sampler/internal/settings",
+        "//agentplane/central_sampler/protocol",
+        "//agentplane/central_sampler/providers/responses",
+        "//agentplane/identity/auth",
+        "@io_k8s_api//authentication/v1:authentication",
+        "@io_k8s_apimachinery//pkg/apis/meta/v1:meta",
+        "@org_golang_google_grpc//codes",
+        "@org_golang_google_grpc//metadata",
+        "@org_golang_google_grpc//status",
+    ],
+)
+
+go_test(
+    name = "integration_test",
+    size = "medium",
+    srcs = [
+        "request_read_service_test.go",
+        "test_helpers_test.go",
+    ],
+    embed = [":rpc"],
+    race = "on",
+    tags = ["requires_docker"],
+    deps = [
+        "//agentplane/central_sampler/api:sampler_go_proto",
+        "//agentplane/central_sampler/internal/sessionstore",
+        "//agentplane/central_sampler/internal/settings",
+        "//agentplane/central_sampler/migrations",
+        "//agentplane/central_sampler/protocol",
+        "//agentplane/central_sampler/providers/responses",
+        "//agentplane/identity/auth",
+        "//agentplane/proto:identity_go_proto",
+        "//agentplane/proto:session_go_proto",
+        "//agentplane/testing/postgres",
+        "@com_github_google_uuid//:uuid",
+        "@com_github_lib_pq//:pq",
+        "@io_k8s_api//authentication/v1:authentication",
+        "@io_k8s_apimachinery//pkg/apis/meta/v1:meta",
+        "@org_golang_google_grpc//:grpc",
+        "@org_golang_google_grpc//codes",
+        "@org_golang_google_grpc//metadata",
+        "@org_golang_google_grpc//status",
+    ],
+)

agentplane/central_sampler/provider_request_read.go → agentplane/central_sampler/internal/rpc/provider_request_read.go

@@ -1,3 +1,3 @@
-package centralsampler
+package rpc
 
 const maxProviderRequestBodyChunkSize = 1 << 20

agentplane/central_sampler/request_read_service_test.go → agentplane/central_sampler/internal/rpc/request_read_service_test.go

@@ -1,21 +1,27 @@
-package centralsampler
+package rpc
 
 import (
 	"bytes"
 	"context"
+	"database/sql"
+	"os"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/agentydragon/ducktape/agentplane/central_sampler/internal/sessionstore"
+	"github.com/agentydragon/ducktape/agentplane/central_sampler/migrations"
 
 	samplerpb "github.com/agentydragon/ducktape/agentplane/central_sampler/api/samplerpb"
 	samplersettings "github.com/agentydragon/ducktape/agentplane/central_sampler/internal/settings"
 	"github.com/agentydragon/ducktape/agentplane/central_sampler/protocol"
 	"github.com/agentydragon/ducktape/agentplane/central_sampler/providers/responses"
 	"github.com/agentydragon/ducktape/agentplane/identity/auth"
+	identitypb "github.com/agentydragon/ducktape/agentplane/proto/identity"
 	sessionpb "github.com/agentydragon/ducktape/agentplane/proto/session"
+	"github.com/agentydragon/ducktape/agentplane/testing/postgres"
 	"github.com/google/uuid"
+	_ "github.com/lib/pq"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -28,6 +34,51 @@ type requestReadStream struct {
 	responses []*samplerpb.ReadProviderRequestResponse
 }
 
+var rpcIntegrationDatabase *postgres.Database
+
+func TestMain(m *testing.M) {
+	migrate := func(ctx context.Context, db *sql.DB) error { return migrations.Up(ctx, db) }
+	os.Exit(postgres.RunIntegrationTests(migrate, func(database *postgres.Database) int {
+		rpcIntegrationDatabase = database
+		return m.Run()
+	}))
+}
+
+func newTestStore(t testing.TB) *sessionstore.Store {
+	t.Helper()
+	dsn := rpcIntegrationDatabase.Clone(t)
+	db, err := sql.Open("postgres", dsn)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = db.Close() })
+	return sessionstore.New(db, 30*time.Second)
+}
+
+func openStoreSession(t *testing.T, store *sessionstore.Store, request *samplerpb.CreateSessionRequest, ownerUID string, api protocol.APIStyle) (*sessionstore.Session, error) {
+	t.Helper()
+	if request.GetExecutionIdentity() == nil {
+		request.ExecutionIdentity = &identitypb.Identity{Name: "fixture-session-identity", Uid: "fixture-session-identity-uid"}
+	}
+	if request.GetExecutor() == nil {
+		request.Executor = &identitypb.SandboxDestination{
+			Owner:  &identitypb.ServiceAccount{Namespace: "fixture", Name: "executor"},
+			Target: &identitypb.EnvironmentTarget{Namespace: "fixture", Name: "fixture-sandbox", Uid: "fixture-sandbox-uid"},
+		}
+	}
+	if request.GetSpec() != nil && request.GetSpec().GetPlacement() == nil {
+		request.Spec.Placement = &sessionpb.SessionPlacement{Kind: &sessionpb.SessionPlacement_CentralSamplerV1{CentralSamplerV1: &sessionpb.CentralSamplerV1Placement{
+			ExecutorEnvironment: &sessionpb.ExactEnvironment{Name: request.GetExecutor().GetTarget().GetName(), Uid: request.GetExecutor().GetTarget().GetUid()},
+		}}}
+	}
+	platformInstructions := "Fixture Agentplane platform guidance."
+	return store.CreateSession(t.Context(), request, ownerUID, api, &platformInstructions)
+}
+
+func providerAttemptTestPolicyPointer() *sessionstore.ProviderAttemptPolicy {
+	return &sessionstore.ProviderAttemptPolicy{Budget: 3, Lease: time.Minute}
+}
+
 func (s *requestReadStream) Context() context.Context { return s.ctx }
 
 func (s *requestReadStream) Send(response *samplerpb.ReadProviderRequestResponse) error {
@@ -54,7 +105,7 @@ func TestReadProviderRequestStreamsAuthorizedContentAndHidesCrossAccountRequests
 
 	config := samplersettings.ModelConfig{
 		Connection: "fixture", ProviderModel: "fixture-model", API: protocol.APIStyleResponses,
-		ExecEncoding: protocol.ExecEncodingJSON, Responses: responses.Config{Dialect: responses.DialectLite},
+		Responses: responses.Config{Dialect: responses.DialectLite}, ExecEncoding: protocol.ExecEncodingJSON,
 	}
 	connection := samplersettings.ConnectionConfig{BaseURL: "https://provider.example.test"}
 	intent := sessionstore.ProviderRequestIntent(sessionstore.PendingIntentSampling, connection, config)
@@ -74,7 +125,9 @@ func TestReadProviderRequestStreamsAuthorizedContentAndHidesCrossAccountRequests
 	if err != nil {
 		t.Fatal(err)
 	}
-	equivalentJSON(t, stored.Body, wire)
+	if !bytes.Equal(stored.Body, wire) {
+		t.Fatal("stored request body differs from prepared request bytes")
+	}
 	if len(stored.Body) <= maxProviderRequestBodyChunkSize {
 		t.Fatalf("fixture body is %d bytes; want more than one chunk", len(stored.Body))
 	}

agentplane/central_sampler/service.go → agentplane/central_sampler/internal/rpc/service.go

@@ -1,4 +1,4 @@
-package centralsampler
+package rpc
 
 import (
 	"context"
@@ -9,6 +9,7 @@ import (
 	"strings"
 	"time"
 
+	centralsampler "github.com/agentydragon/ducktape/agentplane/central_sampler"
 	"github.com/agentydragon/ducktape/agentplane/central_sampler/internal/sessionstore"
 
 	samplerpb "github.com/agentydragon/ducktape/agentplane/central_sampler/api/samplerpb"
@@ -31,7 +32,7 @@ import (
 type Sampler struct {
 	samplerpb.UnimplementedSamplerServer
 	Store    *sessionstore.Store
-	Runtime  *Runtime
+	Runtime  *centralsampler.Runtime
 	Resolver *auth.Resolver
 	Settings samplersettings.Settings
 	Wakeups  *wakeups.Hub

agentplane/central_sampler/service_test.go → agentplane/central_sampler/internal/rpc/service_test.go

@@ -1,7 +1,6 @@
-package centralsampler
+package rpc
 
 import (
-	"context"
 	"errors"
 	"testing"
 
@@ -12,22 +11,8 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
-	authenticationv1 "k8s.io/api/authentication/v1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
-type fixtureReviews struct {
-	uid, account string
-	failure      error
-}
-
-func (r fixtureReviews) Create(ctx context.Context, request *authenticationv1.TokenReview, options metav1.CreateOptions) (*authenticationv1.TokenReview, error) {
-	if r.failure != nil {
-		return nil, r.failure
-	}
-	return &authenticationv1.TokenReview{Status: authenticationv1.TokenReviewStatus{Authenticated: request.Spec.Token == "test-token", Audiences: request.Spec.Audiences, User: authenticationv1.UserInfo{Username: r.account, UID: r.uid, Extra: map[string]authenticationv1.ExtraValue{"authentication.kubernetes.io/pod-name": {"test-pod"}, "authentication.kubernetes.io/pod-uid": {"test-pod-uid"}}}}}, nil
-}
-
 func TestAdmissionAuthenticatesAndAuthorizesAccountUID(t *testing.T) {
 	for _, c := range []struct {
 		name, uid, account string

agentplane/central_sampler/internal/rpc/test_helpers_test.go

@@ -0,0 +1,32 @@
+package rpc
+
+import (
+	"context"
+
+	samplersettings "github.com/agentydragon/ducktape/agentplane/central_sampler/internal/settings"
+	"github.com/agentydragon/ducktape/agentplane/central_sampler/protocol"
+	"github.com/agentydragon/ducktape/agentplane/central_sampler/providers/responses"
+	authenticationv1 "k8s.io/api/authentication/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+type fixtureReviews struct {
+	uid, account string
+	failure      error
+}
+
+func (r fixtureReviews) Create(ctx context.Context, request *authenticationv1.TokenReview, options metav1.CreateOptions) (*authenticationv1.TokenReview, error) {
+	if r.failure != nil {
+		return nil, r.failure
+	}
+	return &authenticationv1.TokenReview{Status: authenticationv1.TokenReviewStatus{Authenticated: request.Spec.Token == "test-token", Audiences: request.Spec.Audiences, User: authenticationv1.UserInfo{Username: r.account, UID: r.uid, Extra: map[string]authenticationv1.ExtraValue{"authentication.kubernetes.io/pod-name": {"test-pod"}, "authentication.kubernetes.io/pod-uid": {"test-pod-uid"}}}}}, nil
+}
+
+func fixtureSettings() samplersettings.Settings {
+	platformInstructions := "Fixture Agentplane guidance."
+	return samplersettings.Settings{
+		Models:      map[string]samplersettings.ModelConfig{"test-model": {Connection: "provider", ProviderModel: "test-model", API: protocol.APIStyleResponses, Responses: responses.Config{Dialect: responses.DialectLite}, ExecEncoding: protocol.ExecEncodingFreeform, ReasoningEfforts: []string{"", "high"}}},
+		Connections: map[string]samplersettings.ConnectionConfig{"provider": {BaseURL: "https://model.invalid", Auth: samplersettings.AuthConfig{Type: samplersettings.AuthTypeBearer, Credential: "test-credential"}}},
+		DatabaseURL: "postgresql://fixture.invalid/sampler", Listen: "127.0.0.1:0", Audience: "test-sampler", AllowedAccounts: []string{"test/client"}, PlatformInstructions: &platformInstructions, CodeModeHost: "/fixture/host", ExecutorCodex: "/fixture/codex", Logs: "/fixture/logs", DirectExecutorTarget: "127.0.0.1:4501", DirectExecutorInsecureLocaltest: true, MaxLiveSessions: 2, RPCTimeoutSeconds: 10, LeaseSeconds: 60,
+	}
+}