Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ func managedArchiveFixture(t *testing.T) (*Store, *Store, string) {
if err != nil {
t.Fatal(err)
}
s := withPlacement(t, NewWithCredentialCipher(pool, cipher))
s := NewWithCredentialCipher(pool, cipher)
w := executionWriter(t, s)
installation := uuid.NewString()
changes := deploymentExecution(t, w)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,11 +36,10 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) {
if err != nil {
t.Fatal(err)
}
s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher}
s.SetPlacement(fixtureRules(t, db))
s := NewWithCredentialCipher(pool, cipher)
installation := uuid.NewString()
provider := &lifecycleProvider{resources: map[string]sandbox.Info{}}
deployments := fixtureDeployment(t, db)
deployments := deploymentService(t, s)
providerConfig := func(setup deployment.Setup) *execution.RuntimeProvider {
return &execution.RuntimeProvider{InstallationID: setup.InstallationID, ProviderKind: setup.Provider, Generation: setup.Generation, Mode: setup.Mode, AdmissionPaused: setup.AdmissionPaused, CoreURL: "https://core.example/api/v1", BackendFingerprint: setup.BackendFingerprint, Provider: provider}
}
Expand All @@ -57,11 +56,11 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) {
if err != nil {
t.Fatal(err)
}
owner, err := fixtureOwner(db, lease)
owner, err := fixtureOwner(s, lease)
if err != nil {
t.Fatal(errors.Join(err, lease.Close(t.Context())))
}
worker := startOwnedWorker(t, t.Context(), db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}, owner)
worker := startOwnedWorker(t, t.Context(), s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}, owner)
var once sync.Once
stop := func() {
once.Do(func() {
Expand All @@ -77,7 +76,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) {
}
projectID := uuid.NewString()
ctx := adminaudit.WithSource(t.Context(), adminaudit.Source{CredentialID: "fixture-admin", ProjectID: projectID, RequestID: uuid.NewString(), TraceID: uuid.NewString()})
_, management := fixtureProjects(t, db)
_, management := fixtureProjects(t, s)
project, err := management.CreateProject(ctx, projects.CreateProject{ID: projectID, Name: "Archive HTTP fixture"})
if err != nil {
t.Fatal(err)
Expand All @@ -103,7 +102,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) {
if err != nil {
t.Fatal(err)
}
handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), workerExecution(t, worker), func(d *api.Dependencies) { d.Execution.SessionArchive = owner.Deployment }, withCoreKeys(admin))
handler, err := publicHandler(t, s, nil, "codex", storeKeys(s), workerExecution(t, worker), func(d *api.Dependencies) { d.Execution.SessionArchive = owner.Deployment }, withCoreKeys(admin))
if err != nil {
t.Fatal(err)
}
Expand All @@ -121,7 +120,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) {
if w.Code != http.StatusOK || json.Unmarshal(w.Body.Bytes(), &archived) != nil || archived.State != "cleanup_pending" || archived.SessionID != active.ID {
t.Fatalf("archive failed: %d %s", w.Code, w.Body)
}
allocation, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: project.TenantID, EnvironmentID: active.Environment.ID})
allocation, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: project.TenantID, EnvironmentID: active.Environment.ID})
if err != nil || allocation.ID != allocated.ID || allocation.State != "cleanup_pending" {
t.Fatal("archive did not retain cleanup ownership", allocation, err)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,12 @@ import (
func TestAgentExecutionDefaultsPublicSnapshotAndPrecedence(t *testing.T) {
_, pool := testStore(t)
cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{19}, 32))
st, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher}
st.SetPlacement(fixtureRules(t, db))
st := NewWithCredentialCipher(pool, cipher)
tenant, token := uuid.NewString(), uuid.NewString()
auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "defaults-test", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}})
deployment := &v1.ModelProviderInput{Protocol: "responses", BaseURL: "https://deployment.example/v1", APIKey: "deployment-canary"}
defaultsCalls := 0
handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, db), withHarnesses([]string{"codex", "claude_sdk", "mcode"}), modelProviderDefaults(func(context.Context, string) (*modelconfiguration.Snapshot, error) {
handler, err := publicHandler(t, st, auth, "codex", storeExecution(t, st), managedSandboxes(t, st), withHarnesses([]string{"codex", "claude_sdk", "mcode"}), modelProviderDefaults(func(context.Context, string) (*modelconfiguration.Snapshot, error) {
defaultsCalls++
copy := *deployment
return &modelconfiguration.Snapshot{Model: "fixture", Provider: &copy, Revision: uuid.New()}, nil
Expand Down
8 changes: 4 additions & 4 deletions services/core/tests/integration/agents_delete_public_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,20 +17,20 @@ func TestAgentDeletionOfficialClient(t *testing.T) {
if python == "" {
t.Skip("pinned official Python SDK required")
}
s, db := newTestStoreDB(t)
s, _ := testStore(t)
token, foreign := uuid.NewString(), uuid.NewString()
auth := newTestAuthenticator(t, []testAPIKey{
{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: uuid.NewString()},
{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()},
})
h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s))
h, err := publicHandler(t, s, auth, "codex", storeExecution(t, s))
if err != nil {
t.Fatal(err)
}
server := httptest.NewServer(h)
defer server.Close()
recoveredStore, recoveredDB := New(db.pool), fixtureDB{pool: db.pool}
h, err = publicHandler(t, recoveredStore, recoveredDB, auth, "codex", storeExecution(t, recoveredStore))
recoveredStore := New(s.pool)
h, err = publicHandler(t, recoveredStore, auth, "codex", storeExecution(t, recoveredStore))
if err != nil {
t.Fatal(err)
}
Expand Down
8 changes: 4 additions & 4 deletions services/core/tests/integration/agents_update_public_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,20 +17,20 @@ func TestAgentUpdateOfficialClient(t *testing.T) {
if python == "" {
t.Skip("pinned official Python SDK required")
}
s, db := newTestStoreDB(t)
s, _ := testStore(t)
token, foreign := uuid.NewString(), uuid.NewString()
auth := newTestAuthenticator(t, []testAPIKey{
{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: uuid.NewString()},
{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()},
})
h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s))
h, err := publicHandler(t, s, auth, "codex", storeExecution(t, s))
if err != nil {
t.Fatal(err)
}
server := httptest.NewServer(h)
defer server.Close()
recoveredStore, recoveredDB := New(db.pool), fixtureDB{pool: db.pool}
h, err = publicHandler(t, recoveredStore, recoveredDB, auth, "codex", storeExecution(t, recoveredStore))
recoveredStore := New(s.pool)
h, err = publicHandler(t, recoveredStore, auth, "codex", storeExecution(t, recoveredStore))
if err != nil {
t.Fatal(err)
}
Expand Down
39 changes: 20 additions & 19 deletions services/core/tests/integration/archive_cancellation_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,8 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
for _, scenario := range []string{"receipt_without_heartbeat", "heartbeat_before_receipt", "done_heartbeat_ack", "ack_commit_blocked", "rotated", "expired", "transport_lost", "negative_ack", "missing_outcome", "revoke_before_archive", "cancel_revoke_archive", "revoke_after_archive", "revoke_concurrent_archive"} {
t.Run(scenario, func(t *testing.T) {
heartbeat := scenario != "receipt_without_heartbeat"
s, db := newManagedTestStoreDB(t)
leased := executionOwner(t, db)
s, _ := newManagedTestStore(t)
leased := executionOwner(t, s)
t.Cleanup(func() {
if err := leased.Lease.Close(context.Background()); err != nil {
t.Error(err)
Expand All @@ -52,7 +52,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
}
projectID := uuid.NewString()
auditCtx := adminaudit.WithSource(t.Context(), adminaudit.Source{CredentialID: "fixture-admin", ProjectID: projectID, RequestID: uuid.NewString(), TraceID: uuid.NewString()})
_, management := fixtureProjects(t, db)
_, management := fixtureProjects(t, s)
project, err := management.CreateProject(auditCtx, projects.CreateProject{ID: projectID, Name: "Archive diagnosis"})
if err != nil {
t.Fatal(err)
Expand Down Expand Up @@ -85,7 +85,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
}
server := httptest.NewUnstartedServer(nil)
wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws"
handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), fixtureSessionStore(db), wsURL)
handler, registry, err := runtime.NewGateway(sessionAdapter(s), sessionService(t, s), sessionAdapter(s), wsURL)
if err != nil {
t.Fatal(err)
}
Expand All @@ -99,11 +99,12 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
t.Fatal(err)
}
t.Cleanup(func() { conn.Close() })
sessionStore, service, err := fixtureSessions(db)
sessionStore := sessionAdapter(s)
service, err := newSessionService(s)
if err != nil {
t.Fatal(err)
}
h := &dispatchHarness{t: t, s: s, db: db, lease: leased.Lease, owned: &leased, tenant: project.TenantID, session: session, conn: conn, registry: registry, d: &execution.Dispatcher{Registry: registry, Observer: modelconfigurationpg.New(pgunit.NewPool(db.pool), db.cipher), Sessions: service, SessionsReader: sessionStore}}
h := &dispatchHarness{t: t, s: s, lease: leased.Lease, owned: &leased, tenant: project.TenantID, session: session, conn: conn, registry: registry, d: &execution.Dispatcher{Registry: registry, Observer: modelconfigurationpg.New(pgunit.NewPool(s.pool), s.credentialCipher), Sessions: service, SessionsReader: sessionStore}}
h.d = h.bound()
capabilities := workerEnvironmentCapabilities()
capabilities.FunctionTools = proto.CapabilitySupported
Expand Down Expand Up @@ -146,14 +147,14 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
}
}
if scenario == "revoke_before_archive" || scenario == "cancel_revoke_archive" {
if err := fixtureSessionService(t, db).RevokeDevice(t.Context(), h.tenant, owner.DeviceID); err != nil {
if err := sessionService(t, s).RevokeDevice(t.Context(), h.tenant, owner.DeviceID); err != nil {
t.Fatal(err)
}
}
var revokeDone chan error
if scenario == "revoke_concurrent_archive" {
revokeDone = make(chan error, 1)
go func() { revokeDone <- fixtureSessionService(t, db).RevokeDevice(t.Context(), h.tenant, owner.DeviceID) }()
go func() { revokeDone <- sessionService(t, s).RevokeDevice(t.Context(), h.tenant, owner.DeviceID) }()
}

archived, err := leased.Deployment.ArchiveSession(auditCtx, h.tenant, session.ID, 1)
Expand All @@ -166,7 +167,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
}
}
if scenario == "revoke_after_archive" {
if err := fixtureSessionService(t, db).RevokeDevice(t.Context(), h.tenant, owner.DeviceID); err != nil {
if err := sessionService(t, s).RevokeDevice(t.Context(), h.tenant, owner.DeviceID); err != nil {
t.Fatal(err)
}
}
Expand All @@ -184,7 +185,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
if err != nil || current.Status != sessions.TurnWaiting || current.CancelRequestedAt.IsZero() {
t.Fatal("archive must request rather than invent cancellation", current, err)
}
if _, err := runtimegateway.NewAuthenticator(fixtureSessionStore(db)).AuthenticateBearer(t.Context(), owner.DeviceID, secret); !errors.Is(err, runtimegateway.ErrAuthUnknownDevice) {
if _, err := runtimegateway.NewAuthenticator(sessionAdapter(s)).AuthenticateBearer(t.Context(), owner.DeviceID, secret); !errors.Is(err, runtimegateway.ErrAuthUnknownDevice) {
t.Fatal("archive allowed renewed authority", err)
}
rejected, response, dialErr := websocket.DefaultDialer.Dial(u.String(), http.Header{"Authorization": {"Bearer " + secret}})
Expand All @@ -197,11 +198,11 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
if dialErr == nil || response == nil || response.StatusCode != http.StatusUnauthorized {
t.Fatal("revoked Runtime reconnected")
}
drain, err := fixtureSessionStore(db).ArchivedCancellationReceipt(t.Context(), owner.DeviceID, secret, nil)
drain, err := sessionAdapter(s).ArchivedCancellationReceipt(t.Context(), owner.DeviceID, secret, nil)
if err != nil || drain.RunID != "" {
t.Fatal("unowned delivery got receipt permission", drain, err)
}
drain, err = fixtureSessionStore(db).ArchivedCancellationReceipt(t.Context(), owner.DeviceID, runtimedevice.HashCredential(secret), []string{input.TurnID})
drain, err = sessionAdapter(s).ArchivedCancellationReceipt(t.Context(), owner.DeviceID, runtimedevice.HashCredential(secret), []string{input.TurnID})
if err != nil || (drain.RunID == input.TurnID) == strings.Contains(scenario, "revoke") {
t.Fatal("archive revocation causality lost", drain, err)
}
Expand All @@ -218,18 +219,18 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
t.Fatal("missing cancel delivery identity")
}
if scenario == "rotated" {
if _, err := db.pool.Exec(t.Context(), "UPDATE devices SET credential_hash=$2 WHERE id=$1", owner.DeviceID, runtimedevice.HashCredential(uuid.NewString())); err != nil {
if _, err := s.pool.Exec(t.Context(), "UPDATE devices SET credential_hash=$2 WHERE id=$1", owner.DeviceID, runtimedevice.HashCredential(uuid.NewString())); err != nil {
t.Fatal(err)
}
}
if scenario == "expired" {
if _, err := db.pool.Exec(t.Context(), "UPDATE turns SET cancel_requested_at=clock_timestamp()-interval '21 seconds' WHERE id=$1", input.TurnID); err != nil {
if _, err := s.pool.Exec(t.Context(), "UPDATE turns SET cancel_requested_at=clock_timestamp()-interval '21 seconds' WHERE id=$1", input.TurnID); err != nil {
t.Fatal(err)
}
}
var unlockCommit func()
if scenario == "ack_commit_blocked" {
tx, err := db.pool.Begin(t.Context())
tx, err := s.pool.Begin(t.Context())
if err != nil {
t.Fatal(err)
}
Expand Down Expand Up @@ -269,7 +270,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
// Observe actual SQL lock contention, not an assumed timing delay.
for deadline := time.Now().Add(3 * time.Second); ; {
var blocked bool
if err := db.pool.QueryRow(t.Context(), "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock' AND query ILIKE '%session_devices%')").Scan(&blocked); err != nil {
if err := s.pool.QueryRow(t.Context(), "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock' AND query ILIKE '%session_devices%')").Scan(&blocked); err != nil {
t.Fatal(err)
}
if blocked {
Expand Down Expand Up @@ -302,7 +303,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
}

var receipts int
if err := db.pool.QueryRow(t.Context(), "SELECT count(*) FROM turn_events WHERE turn_id=$1 AND kind='cancel_receipt'", input.TurnID).Scan(&receipts); err != nil {
if err := s.pool.QueryRow(t.Context(), "SELECT count(*) FROM turn_events WHERE turn_id=$1 AND kind='cancel_receipt'", input.TurnID).Scan(&receipts); err != nil {
t.Fatal(err)
}
wantReceipts := 1
Expand All @@ -314,12 +315,12 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) {
}
// The original cleanup owner survives every delivery outcome; only
// provider receipts can release its resources.
allocation, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: h.tenant, EnvironmentID: session.Environment.ID})
allocation, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: h.tenant, EnvironmentID: session.Environment.ID})
if err != nil || allocation.State != "cleanup_pending" {
t.Fatal(allocation, err)
}
var revoked bool
if err := db.pool.QueryRow(t.Context(), "SELECT revoked_at IS NOT NULL FROM devices WHERE id=$1", owner.DeviceID).Scan(&revoked); err != nil || !revoked {
if err := s.pool.QueryRow(t.Context(), "SELECT revoked_at IS NOT NULL FROM devices WHERE id=$1", owner.DeviceID).Scan(&revoked); err != nil || !revoked {
t.Fatal(revoked, err)
}
})
Expand Down
10 changes: 5 additions & 5 deletions services/core/tests/integration/claude_execution_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ func claudeSession(t *testing.T, h *dispatchHarness, configuration string, prebo
t.Fatal(err)
}
if prebound {
if err := bindSessionDevice(t, h.db, h.tenant, h.session.ID, h.device.ID); err != nil {
if err := bindSessionDevice(t, h.s, h.tenant, h.session.ID, h.device.ID); err != nil {
t.Fatal(err)
}
}
Expand Down Expand Up @@ -54,7 +54,7 @@ func TestClaudeWorkerSelectsStoredEngineAndRestrictiveCapabilities(t *testing.T)
input := h.message("start", "Look up ticket")
ctx, cancel := context.WithCancel(t.Context())
defer cancel()
worker := startWorker(t, ctx, h.db, h.d)
worker := startWorker(t, ctx, h.s, h.d)
done := make(chan error, 1)
go func() { done <- worker.Run(ctx) }()
defer func() {
Expand All @@ -72,7 +72,7 @@ func TestClaudeWorkerSelectsStoredEngineAndRestrictiveCapabilities(t *testing.T)
t.Fatal(turn, err)
}
if !prebound {
if _, err := fixtureSessionStore(h.db).GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) {
if _, err := sessionAdapter(h.s).GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) {
t.Fatal("bound an incapable device", err)
}
}
Expand Down Expand Up @@ -122,7 +122,7 @@ func TestClaudeDispatcherRejectsUnsupportedConfigurationBeforeClaim(t *testing.T
func TestClaudeInvalidImageResultRejectsWholeBatchBeforePersistence(t *testing.T) {
h := newDispatchHarness(t)
claudeSession(t, h, functionConfiguration, false)
worker := startOwnedWorker(t, t.Context(), h.db, h.d, h.owner())
worker := startOwnedWorker(t, t.Context(), h.s, h.d, h.owner())
defer func() { ctx, cancel := context.WithCancel(context.Background()); cancel(); _ = worker.Run(ctx) }()
input := h.message("start", "Run")
if _, err := transitionTurn(t.Context(), h.s, h.tenant, h.session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil {
Expand All @@ -143,7 +143,7 @@ func TestClaudeInvalidImageResultRejectsWholeBatchBeforePersistence(t *testing.T
if _, err := worker.SubmitInputs(t.Context(), h.tenant, h.session.ID, "batch", batch); !errors.Is(err, sessions.ErrInvalidInput) {
t.Fatal(err)
}
saved, err := FixtureFunctionCall(t.Context(), h.db.pool, h.tenant, h.session.ID, input.TurnID, call.CallID)
saved, err := FixtureFunctionCall(t.Context(), h.s.pool, h.tenant, h.session.ID, input.TurnID, call.CallID)
if err != nil || saved.Result != nil || saved.Applied {
t.Fatal(saved, err)
}
Expand Down
4 changes: 2 additions & 2 deletions services/core/tests/integration/claude_mcp_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ func TestClaudeMCPWaitsForCapableRuntime(t *testing.T) {
}
ctx, cancel := context.WithCancel(t.Context())
defer cancel()
worker := startOwnedWorker(t, ctx, h.db, h.d, h.owner())
worker := startOwnedWorker(t, ctx, h.s, h.d, h.owner())
done := make(chan error, 1)
go func() { done <- worker.Run(ctx) }()
defer func() {
Expand All @@ -60,7 +60,7 @@ func TestClaudeMCPWaitsForCapableRuntime(t *testing.T) {
t.Fatal("incapable runtime claimed work", turn, err)
}
if !prebound {
if _, err := fixtureSessionStore(h.db).GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) {
if _, err := sessionAdapter(h.s).GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) {
t.Fatal("bound an incapable runtime", err)
}
}
Expand Down
Loading
Loading