From f9687d3f49e23a7a04f51dc40dd99ac9173ba744 Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Tue, 6 Oct 2026 22:45:19 +0000 Subject: [PATCH] Fold the integration fixture database into Store --- .../integration/admin_session_archive_test.go | 2 +- .../admin_session_archive_worker_http_test.go | 15 +- .../agent_execution_defaults_http_test.go | 5 +- .../integration/agents_delete_public_test.go | 8 +- .../integration/agents_update_public_test.go | 8 +- .../integration/archive_cancellation_test.go | 39 ++--- .../integration/claude_execution_test.go | 10 +- .../core/tests/integration/claude_mcp_test.go | 4 +- .../tests/integration/command_output_test.go | 8 +- .../configuration_validation_public_test.go | 6 +- .../creation_stream_settlement_public_test.go | 8 +- .../credential_matrix_http_test.go | 19 ++- .../integration/deployment_fixture_test.go | 16 +- .../deployment_model_providers_http_test.go | 33 ++-- .../deployment_public_fixture_test.go | 49 ++---- .../core/tests/integration/dispatch_test.go | 30 ++-- .../integration/environment_admission_test.go | 2 +- .../environment_claim_worker_test.go | 14 +- .../environment_connection_events_test.go | 5 +- .../environment_connection_worker_test.go | 14 +- .../integration/environment_device_test.go | 8 +- .../integration/environment_directory_test.go | 4 +- .../environment_executor_command_test.go | 8 +- .../environment_expiry_dispatch_test.go | 4 +- .../environment_expiry_worker_test.go | 18 +-- ...onment_file_write_semantics_public_test.go | 6 +- .../environment_initial_public_test.go | 8 +- .../environment_initialization_test.go | 43 +++--- .../environment_mcp_public_test.go | 6 +- .../integration/environment_plugins_test.go | 2 +- .../environment_retrieve_public_test.go | 22 +-- .../environment_runtime_fixture_test.go | 16 +- .../integration/environment_setup_test.go | 2 +- .../integration/environment_work_test.go | 2 +- .../environment_worker_scan_test.go | 4 +- .../integration/environment_worker_test.go | 6 +- .../file_resource_semantics_public_test.go | 4 +- .../core/tests/integration/fixture_db_test.go | 141 ------------------ .../core/tests/integration/fixtures_test.go | 2 +- .../function_execution_native_test.go | 6 +- .../integration/function_execution_test.go | 8 +- .../function_images_native_test.go | 6 +- .../function_input_execution_test.go | 2 +- .../function_inputs_public_test.go | 8 +- .../function_public_native_test.go | 6 +- .../integration/function_state_public_test.go | 6 +- .../function_stream_native_test.go | 2 +- .../tests/integration/function_worker_test.go | 10 +- .../integration/harness_onboarding_test.go | 4 +- ...sted_initialization_failure_public_test.go | 38 +++-- .../integration/initial_files_http_test.go | 7 +- .../input_conflicts_public_test.go | 18 +-- .../integration/item_order_migration_test.go | 2 +- .../core/tests/integration/item_order_test.go | 8 +- .../core/tests/integration/item_reads_test.go | 32 ++-- .../integration/list_cursor_public_test.go | 11 +- .../integration/list_query_public_test.go | 6 +- .../integration/local_artifact_export_test.go | 2 +- .../local_environment_file_write_test.go | 8 +- .../local_environment_worker_test.go | 10 +- .../tests/integration/managed_fixture_test.go | 2 +- .../integration/mcode_public_native_test.go | 4 +- .../mcp_credential_selection_public_test.go | 4 +- .../message_image_admission_test.go | 2 +- .../integration/message_images_native_test.go | 6 +- .../integration/model_protocol_native_test.go | 6 +- .../integration/native_environment_test.go | 4 +- .../native_public_execution_test.go | 4 +- .../tests/integration/no_environment_test.go | 2 +- .../path_id_semantics_public_test.go | 6 +- .../integration/prepared_dispatch_test.go | 2 +- .../integration/project_api_keys_http_test.go | 4 +- .../projects_public_fixture_test.go | 6 +- .../integration/public_execution_test.go | 12 +- .../public_handler_fixture_test.go | 57 ++++--- .../remote_mcp_credentials_test.go | 12 +- .../core/tests/integration/remote_mcp_test.go | 18 +-- .../integration/request_body_public_test.go | 8 +- .../tests/integration/root_fixture_test.go | 8 +- .../runtime_capabilities_pending_test.go | 13 +- .../runtime_compute_lifecycle_test.go | 31 ++-- .../runtime_configuration_cleanup_test.go | 14 +- .../integration/runtime_connection_test.go | 16 +- .../runtime_creation_settlement_test.go | 26 ++-- .../runtime_deployment_worker_test.go | 14 +- .../runtime_enrollment_connection_test.go | 28 ++-- .../runtime_initialization_test.go | 31 ++-- .../runtime_input_admission_test.go | 6 +- .../integration/runtime_lifecycle_test.go | 76 +++++----- .../runtime_node_lifecycle_fixture_test.go | 16 +- .../runtime_node_lifecycle_test.go | 14 +- .../runtime_observation_scan_test.go | 14 +- .../tests/integration/runtime_pending_test.go | 16 +- .../tests/integration/runtime_scan_test.go | 16 +- .../runtime_wake_hint_integration_test.go | 8 +- .../runtime_worker_recovery_test.go | 6 +- .../sandbox_deployment_resources_test.go | 2 +- .../sandbox_deployment_switch_test.go | 12 +- .../sandbox_deployment_switch_worker_test.go | 11 +- .../sandbox_deployment_view_test.go | 2 +- .../sandbox_deployment_worker_test.go | 10 +- .../sandbox_node_auth_order_http_test.go | 8 +- .../sandbox_specification_lifecycle_test.go | 2 +- .../saved_web_search_public_test.go | 10 +- .../self_hosted_cancel_public_test.go | 22 +-- .../self_hosted_initial_public_test.go | 46 +++--- .../session_agent_filter_public_test.go | 8 +- .../session_artifacts_public_test.go | 17 ++- .../session_deletion_execution_test.go | 6 +- .../session_deletion_lifecycle_public_test.go | 18 +-- .../session_deletion_public_test.go | 6 +- .../session_devices_fixture_test.go | 6 +- .../tests/integration/session_events_test.go | 6 +- .../session_execution_configuration_test.go | 4 +- .../session_initial_public_test.go | 8 +- .../session_model_execution_http_test.go | 5 +- .../integration/session_model_options_test.go | 2 +- .../integration/session_reads_fixture_test.go | 7 +- .../session_reference_retry_public_test.go | 12 +- .../integration/sessions_fixture_test.go | 47 ------ .../core/tests/integration/sessions_test.go | 2 +- .../skill_selectors_public_test.go | 8 +- .../skill_version_deletion_public_test.go | 4 +- .../skill_version_deletion_test.go | 2 +- .../tests/integration/skills_public_test.go | 8 +- .../source_files_errors_public_test.go | 4 +- .../integration/stream_authority_http_test.go | 6 +- .../structured_output_dispatch_test.go | 2 +- .../structured_output_native_test.go | 6 +- .../integration/subagent_dispatch_test.go | 2 +- .../subagent_visibility_public_test.go | 8 +- .../template_composition_public_test.go | 7 +- .../template_null_selection_public_test.go | 7 +- .../token_usage_integration_test.go | 10 +- .../integration/tool_policy_native_test.go | 4 +- .../integration/tool_search_native_test.go | 6 +- .../tests/integration/turn_events_test.go | 2 +- .../unified_model_configuration_http_test.go | 4 +- .../unstorable_text_public_test.go | 4 +- .../tests/integration/vaults_fixture_test.go | 8 +- .../whitespace_input_public_test.go | 18 +-- .../tests/integration/worker_capacity_test.go | 4 +- .../tests/integration/worker_fixture_test.go | 112 ++++++++++++++ .../integration/worker_input_race_test.go | 2 +- .../integration/worker_lease_loss_test.go | 4 +- .../worker_preparation_failure_test.go | 6 +- .../tests/integration/worker_wakeup_test.go | 4 +- 147 files changed, 856 insertions(+), 988 deletions(-) delete mode 100644 services/core/tests/integration/fixture_db_test.go delete mode 100644 services/core/tests/integration/sessions_fixture_test.go create mode 100644 services/core/tests/integration/worker_fixture_test.go diff --git a/services/core/tests/integration/admin_session_archive_test.go b/services/core/tests/integration/admin_session_archive_test.go index 85e4024e9..df75725eb 100644 --- a/services/core/tests/integration/admin_session_archive_test.go +++ b/services/core/tests/integration/admin_session_archive_test.go @@ -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) diff --git a/services/core/tests/integration/admin_session_archive_worker_http_test.go b/services/core/tests/integration/admin_session_archive_worker_http_test.go index 6f586d67b..134fd92d0 100644 --- a/services/core/tests/integration/admin_session_archive_worker_http_test.go +++ b/services/core/tests/integration/admin_session_archive_worker_http_test.go @@ -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} } @@ -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() { @@ -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) @@ -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) } @@ -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) } diff --git a/services/core/tests/integration/agent_execution_defaults_http_test.go b/services/core/tests/integration/agent_execution_defaults_http_test.go index eeff542fc..28f587c6c 100644 --- a/services/core/tests/integration/agent_execution_defaults_http_test.go +++ b/services/core/tests/integration/agent_execution_defaults_http_test.go @@ -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: ©, Revision: uuid.New()}, nil diff --git a/services/core/tests/integration/agents_delete_public_test.go b/services/core/tests/integration/agents_delete_public_test.go index 87ea12a0e..0f37041d4 100644 --- a/services/core/tests/integration/agents_delete_public_test.go +++ b/services/core/tests/integration/agents_delete_public_test.go @@ -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) } diff --git a/services/core/tests/integration/agents_update_public_test.go b/services/core/tests/integration/agents_update_public_test.go index c6f49158f..651a044b1 100644 --- a/services/core/tests/integration/agents_update_public_test.go +++ b/services/core/tests/integration/agents_update_public_test.go @@ -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) } diff --git a/services/core/tests/integration/archive_cancellation_test.go b/services/core/tests/integration/archive_cancellation_test.go index d5aa6bf2b..0b64fef50 100644 --- a/services/core/tests/integration/archive_cancellation_test.go +++ b/services/core/tests/integration/archive_cancellation_test.go @@ -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) @@ -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) @@ -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) } @@ -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 @@ -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) @@ -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) } } @@ -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}}) @@ -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) } @@ -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) } @@ -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 { @@ -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 @@ -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) } }) diff --git a/services/core/tests/integration/claude_execution_test.go b/services/core/tests/integration/claude_execution_test.go index e298a94a4..648969f7c 100644 --- a/services/core/tests/integration/claude_execution_test.go +++ b/services/core/tests/integration/claude_execution_test.go @@ -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) } } @@ -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() { @@ -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) } } @@ -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 { @@ -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) } diff --git a/services/core/tests/integration/claude_mcp_test.go b/services/core/tests/integration/claude_mcp_test.go index f9d9f7a0e..32a18777a 100644 --- a/services/core/tests/integration/claude_mcp_test.go +++ b/services/core/tests/integration/claude_mcp_test.go @@ -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() { @@ -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) } } diff --git a/services/core/tests/integration/command_output_test.go b/services/core/tests/integration/command_output_test.go index ca65f2527..01406737e 100644 --- a/services/core/tests/integration/command_output_test.go +++ b/services/core/tests/integration/command_output_test.go @@ -15,7 +15,7 @@ import ( func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { ctx := context.Background() s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "command-output"}) if err != nil { @@ -53,7 +53,7 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { if before != after { t.Fatal("rollback published output") } - page, err := sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) + page, err := sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil || len(page.Items) != 2 { t.Fatalf("read draft: %+v %v", page, err) } @@ -75,7 +75,7 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { // Reopening the Store recovers committed Items without creating events. reopened := New(pool) before, _ = sessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) - page, err = sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) + page, err = sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil || len(page.Items) != 3 { t.Fatalf("recovery: %+v %v", page, err) } @@ -144,7 +144,7 @@ func TestExecutionJournalsCommandOutputBeforeCancellation(t *testing.T) { } h.write(input.TurnID, proto.TypeInteractionDecisionAck, proto.InteractionDecisionAckPayload{DeliveryID: cancel.DeliveryID, Applied: true, Outcome: &proto.DonePayload{}}) h.finished(result, sessions.TurnCancelled) - page, err := sessionReads(h.db.pool).ListItems(ctx, h.tenant, h.session.ID, "", 100, true) + page, err := sessionAdapter(h.s).ListItems(ctx, h.tenant, h.session.ID, "", 100, true) if err != nil || len(page.Items) != 2 || page.Items[1].Status != "incomplete" || page.Items[1].Output != "partial" { t.Fatalf("journal/cancellation lost partial output: %+v %v", page, err) } diff --git a/services/core/tests/integration/configuration_validation_public_test.go b/services/core/tests/integration/configuration_validation_public_test.go index c3ef553c8..e1bbe8f26 100644 --- a/services/core/tests/integration/configuration_validation_public_test.go +++ b/services/core/tests/integration/configuration_validation_public_test.go @@ -24,13 +24,13 @@ func TestAgentConfigurationValidationRejectsWithoutWritesPostgres(t *testing.T) if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) owner, foreign, ownerTenant := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "config-owner", TokenSHA256: runtimedevice.HashCredential(owner), TenantID: ownerTenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "config-foreign", 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) } @@ -43,7 +43,7 @@ func TestAgentConfigurationValidationRejectsWithoutWritesPostgres(t *testing.T) saved := func(tools string) string { return `{"model":"config-model","name":null,"instructions":null,"multi_agent":{"enabled":false,"max_concurrent_subagents":null},"reasoning":{},"service_tier":"auto","text":{"format":{"type":"text"},"verbosity":"medium"},"tools":` + tools + `}` } - _, agentService := fixtureAgents(t, db) + _, agentService := fixtureAgents(t, s) legacy, err := agentService.Create(t.Context(), agents.CreateCommand{TenantID: ownerTenant, Metadata: map[string]string{}, Configuration: json.RawMessage(saved(`[{"type":"function","name":"lookup","description":"","parameters":{"type":"string"},"defer_loading":false}]`))}) if err != nil { t.Fatal(err) diff --git a/services/core/tests/integration/creation_stream_settlement_public_test.go b/services/core/tests/integration/creation_stream_settlement_public_test.go index 98f6043fd..63e64c863 100644 --- a/services/core/tests/integration/creation_stream_settlement_public_test.go +++ b/services/core/tests/integration/creation_stream_settlement_public_test.go @@ -113,10 +113,10 @@ func (s sseLines) open(t *testing.T) { // creation stream whose initial reservation is cancelled without a Session event // ends through the committed projection, while GET stays open. func TestCreationStreamPublicLifetimes(t *testing.T) { - s, db := newModelTestStoreDB(t) + s, _ := NewModelTestStore(t) tenant, token := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}}) - handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://offline-executor.example")) + handler, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), executorURL("https://offline-executor.example")) if err != nil { t.Fatal(err) } @@ -127,7 +127,7 @@ func TestCreationStreamPublicLifetimes(t *testing.T) { handler.ServeHTTP(w, r) })) defer server.Close() - sessionExecution := executionOwner(t, db).Sessions + sessionExecution := executionOwner(t, s).Sessions connect := func(environment string) { t.Helper() generation := uuid.NewString() @@ -203,7 +203,7 @@ func TestCreationStreamPublicLifetimes(t *testing.T) { } fresh.open(t) var reservation string - if err := db.pool.QueryRow(t.Context(), "SELECT id FROM environment_input_reservations WHERE session_id=$1 AND is_initial", session).Scan(&reservation); err != nil { + if err := s.pool.QueryRow(t.Context(), "SELECT id FROM environment_input_reservations WHERE session_id=$1 AND is_initial", session).Scan(&reservation); err != nil { t.Fatal(err) } cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session) diff --git a/services/core/tests/integration/credential_matrix_http_test.go b/services/core/tests/integration/credential_matrix_http_test.go index 395828d20..d424fc7d4 100644 --- a/services/core/tests/integration/credential_matrix_http_test.go +++ b/services/core/tests/integration/credential_matrix_http_test.go @@ -28,23 +28,22 @@ import ( // the Core key in /core/v1, and node and executor credentials only on their // own /api/v1 machine connection routes. func TestCredentialNamespaceMatrix(t *testing.T) { - s, db := newManagedTestStoreDB(t) - db.publicURL = "https://core.example" - s.SetPlacement(fixtureRules(t, db)) + s, _ := newManagedTestStore(t) + s.SetPlacement(placementRules(t, "https://core.example")) ctx := t.Context() coreKey := uuid.NewString() admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential(coreKey)}) if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, db), withCoreKeys(admin)) + handler, err := publicHandler(t, s, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, s), withCoreKeys(admin)) if err != nil { t.Fatal(err) } // The server composition: daemon transport beside the API handler. mux := http.NewServeMux() - mux.Handle("/api/v1/agent-daemon/enroll", runtimeenrollment.EnrollmentHandler(fixtureSessionService(t, db))) - mux.Handle("/api/v1/agent-daemon/connection", runtimeenrollment.ConnectionHandler(fixtureSessionStore(db), runtimegateway.NewRegistry())) + mux.Handle("/api/v1/agent-daemon/enroll", runtimeenrollment.EnrollmentHandler(sessionService(t, s))) + mux.Handle("/api/v1/agent-daemon/connection", runtimeenrollment.ConnectionHandler(sessionAdapter(s), runtimegateway.NewRegistry())) mux.Handle("/", handler) server := api.CanonicalPaths(mux) call := func(method, path, token, body string) *httptest.ResponseRecorder { @@ -71,7 +70,7 @@ func TestCredentialNamespaceMatrix(t *testing.T) { created("POST", "/core/v1/projects/"+project.ID+"/keys", coreKey, `{"name":"application"}`, &projectKey) // An executor credential for a self_hosted Session of that Project. - binding, err := projectpg.New(pgunit.NewPool(db.pool)).GetProject(ctx, project.ID) + binding, err := projectpg.New(pgunit.NewPool(s.pool)).GetProject(ctx, project.ID) if err != nil { t.Fatal(err) } @@ -80,7 +79,7 @@ func TestCredentialNamespaceMatrix(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := fixtureSessionStore(db).GetSessionEnvironment(ctx, binding.Principal.TenantID, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(ctx, binding.Principal.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -88,7 +87,7 @@ func TestCredentialNamespaceMatrix(t *testing.T) { created("POST", "/core/v1/projects/"+project.ID+"/environments/"+environment.ID+"/executor-credentials", coreKey, `{"key_id":"`+uuid.NewString()+`"}`, &executor) // A node credential: a Docker deployment, an enrollment token issued with the Core key, and an enrolled node. - deployments := fixtureDeployment(t, db) + deployments := deploymentService(t, s) installation := uuid.NewString() provider := &lifecycleProvider{resources: map[string]sandbox.Info{}} runtimes := execution.NewDeferredRuntimeProvider(installation, func(ctx context.Context) (*execution.RuntimeProvider, error) { @@ -100,7 +99,7 @@ func TestCredentialNamespaceMatrix(t *testing.T) { }, func(_ context.Context, setup deployment.Setup) (execution.PreparedRuntimeDeployment, error) { return execution.PreparedRuntimeDeployment{Config: &execution.RuntimeProvider{InstallationID: setup.InstallationID, ProviderKind: setup.Provider, Mode: setup.Mode, AdmissionPaused: setup.AdmissionPaused, CoreURL: "https://core.example/api/v1", BackendFingerprint: setup.BackendFingerprint, Provider: provider}}, nil }) - worker := startWorker(t, ctx, db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: runtimes}) + worker := startWorker(t, ctx, s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: runtimes}) var stop sync.Once t.Cleanup(func() { stop.Do(func() { diff --git a/services/core/tests/integration/deployment_fixture_test.go b/services/core/tests/integration/deployment_fixture_test.go index 5b35597c9..25018ae00 100644 --- a/services/core/tests/integration/deployment_fixture_test.go +++ b/services/core/tests/integration/deployment_fixture_test.go @@ -27,20 +27,10 @@ func placementRules(t testing.TB, publicURL string) *placement.Rules { return rules } -// withPlacement gives s the placement rules cmd/server builds on the built-in -// providers and an unset public URL, and returns s. -func withPlacement(t testing.TB, s *Store) *Store { - t.Helper() - s.SetPlacement(placementRules(t, "")) - return s -} - -// deploymentService builds the deployment service as cmd/server does, on s's -// database, credential key and placement rules. +// deploymentService is fixtureDeploymentService(s) for a test. func deploymentService(t testing.TB, s *Store) *deployment.Service { t.Helper() - adapter := deploymentStore(s) - service, err := deployment.NewService(adapter, adapter, providers.Builtin(), s.placement) + service, err := fixtureDeploymentService(s) if err != nil { t.Fatal(err) } @@ -54,7 +44,7 @@ func deploymentExecution(t testing.TB, w *Store) *deployment.ExecutionOperations if w.lease == nil { t.Fatal("deployment execution operations need an execution writer") } - operations, err := deployment.NewExecutionOperations(deploymentService(t, w), deploymentpg.NewExecution(w.lease, w.credentialCipher)) + _, operations, err := fixtureDeploymentExecution(w, w.lease) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/deployment_model_providers_http_test.go b/services/core/tests/integration/deployment_model_providers_http_test.go index a9c561762..7e25421a1 100644 --- a/services/core/tests/integration/deployment_model_providers_http_test.go +++ b/services/core/tests/integration/deployment_model_providers_http_test.go @@ -34,19 +34,18 @@ func TestDeploymentModelProvidersHTTP(t *testing.T) { t.Fatal(err) } cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{53}, 32)) - st, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} - st.SetPlacement(fixtureRules(t, db)) + st := NewWithCredentialCipher(pool, cipher) tenant, projectKey, coreKey := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "defaults-http", TokenSHA256: runtimedevice.HashCredential(projectKey), TenantID: tenant}}) admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential(coreKey)}) if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, db), withCoreKeys(admin), withHarnesses([]string{"codex", "mcode"})) + handler, err := publicHandler(t, st, auth, "codex", storeExecution(t, st), managedSandboxes(t, st), withCoreKeys(admin), withHarnesses([]string{"codex", "mcode"})) if err != nil { t.Fatal(err) } - defaults := deploymentDefaults(t, db) + defaults := deploymentDefaults(t, st) call := func(method, path, key, body string, status int) map[string]json.RawMessage { t.Helper() r := httptest.NewRequest(method, path, strings.NewReader(body)) @@ -247,7 +246,7 @@ func TestLegacySessionWithoutProviderCannotStartWork(t *testing.T) { if err != nil { t.Fatal(err) } - worker, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + worker, stop := startEnvironmentExpiryWorker(t, h.s, h.d) defer stop() _, pool := testStore(t) reservations := func() int { @@ -290,18 +289,18 @@ func TestLegacySessionWithoutProviderCannotStartWork(t *testing.T) { // recorded first: a same-key retry returns the committed Session after the // default was replaced or removed. func TestNoneSessionRetryAfterDeploymentDefaultChanges(t *testing.T) { - st, db := newModelTestStoreDB(t) - if _, err := db.pool.Exec(t.Context(), "DELETE FROM deployment_model_providers"); err != nil { + st, _ := NewModelTestStore(t) + if _, err := st.pool.Exec(t.Context(), "DELETE FROM deployment_model_providers"); err != nil { t.Fatal(err) } tenant, token := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "none-retry", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}}) - handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st)) + handler, err := publicHandler(t, st, auth, "codex", storeExecution(t, st)) if err != nil { t.Fatal(err) } admin := adminaudit.WithSource(t.Context(), adminaudit.Source{CredentialID: "abcd1234", RequestID: "none-retry", TraceID: "none-retry"}) - defaults := deploymentDefaults(t, db) + defaults := deploymentDefaults(t, st) setDefault := func(key string) { t.Helper() if _, err := defaults.Replace(admin, modelconfiguration.Replacement{Harness: "codex", Configuration: v1.ModelConfigurationInput{ModelProvider: v1.ModelProviderInput{Protocol: "responses", BaseURL: "https://deployment.example/v1", APIKey: key}, Model: "fixture"}}); err != nil { @@ -351,18 +350,18 @@ func TestNoneSessionRetryAfterDeploymentDefaultChanges(t *testing.T) { t.Fatal("retry changed the frozen provider", err) } var revision uuid.UUID - if err := db.pool.QueryRow(t.Context(), "SELECT deployment_provider_revision FROM session_execution_configuration WHERE session_id=$1", original).Scan(&revision); err != nil || revision != snapshot.Revision { + if err := st.pool.QueryRow(t.Context(), "SELECT deployment_provider_revision FROM session_execution_configuration WHERE session_id=$1", original).Scan(&revision); err != nil || revision != snapshot.Revision { t.Fatal("API retry changed frozen revision", err) } } func TestDeploymentProviderResolutionPairsRevisionDuringReplacement(t *testing.T) { - st, db := newManagedTestStoreDB(t) + st, _ := newManagedTestStore(t) tenant, token := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "tuple-test", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}}) admin := adminaudit.WithSource(t.Context(), adminaudit.Source{CredentialID: "fixture-admin", RequestID: uuid.NewString(), TraceID: uuid.NewString()}) provider := v1.ModelProviderInput{Protocol: "responses", BaseURL: "https://original.example/v1", APIKey: "original-fixture-key"} - defaults := deploymentDefaults(t, db) + defaults := deploymentDefaults(t, st) if _, err := defaults.Replace(admin, modelconfiguration.Replacement{Harness: "codex", Configuration: v1.ModelConfigurationInput{ModelProvider: provider, Model: "fixture"}}); err != nil { t.Fatal(err) } @@ -382,7 +381,7 @@ func TestDeploymentProviderResolutionPairsRevisionDuringReplacement(t *testing.T _, err = defaults.Replace(admin, modelconfiguration.Replacement{Harness: harness, Configuration: v1.ModelConfigurationInput{ModelProvider: replacement, Model: "fixture"}}) return snapshot, err } - handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), modelProviderDefaults(resolver)) + handler, err := publicHandler(t, st, auth, "codex", storeExecution(t, st), modelProviderDefaults(resolver)) if err != nil { t.Fatal(err) } @@ -397,7 +396,7 @@ func TestDeploymentProviderResolutionPairsRevisionDuringReplacement(t *testing.T t.Fatalf("creation failed: %d %s", w.Code, w.Body) } var revision uuid.UUID - if err = db.pool.QueryRow(t.Context(), "SELECT deployment_provider_revision FROM session_execution_configuration WHERE session_id=$1", session.ID).Scan(&revision); err != nil || revision != original.Revision { + if err = st.pool.QueryRow(t.Context(), "SELECT deployment_provider_revision FROM session_execution_configuration WHERE session_id=$1", session.ID).Scan(&revision); err != nil || revision != original.Revision { t.Fatal("tuple revision changed", err) } frozen, err := sessionAdapter(st).SessionModelExecution(t.Context(), tenant, session.ID) @@ -430,11 +429,11 @@ func TestDeploymentProviderResolutionFixtureIsolation(t *testing.T) { } } -// deploymentDefaults serves deployment default model configurations from db, +// deploymentDefaults serves deployment default model configurations from s, // as the Core routes do. -func deploymentDefaults(t *testing.T, db fixtureDB) *modelconfiguration.Service { +func deploymentDefaults(t *testing.T, s *Store) *modelconfiguration.Service { t.Helper() - service, err := modelconfiguration.NewService(modelconfigurationpg.New(pgunit.NewPool(db.pool), db.cipher)) + service, err := modelconfiguration.NewService(modelconfigurationpg.New(pgunit.NewPool(s.pool), s.credentialCipher)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/deployment_public_fixture_test.go b/services/core/tests/integration/deployment_public_fixture_test.go index dff1550b7..b1ce77186 100644 --- a/services/core/tests/integration/deployment_public_fixture_test.go +++ b/services/core/tests/integration/deployment_public_fixture_test.go @@ -4,52 +4,23 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" ) -// fixtureDeployment builds the deployment service on db, as cmd/server does. -func fixtureDeployment(t testing.TB, db fixtureDB) *deployment.Service { - t.Helper() - service, err := fixtureDeploymentService(db) - if err != nil { - t.Fatal(err) - } - return service -} - -func fixtureDeploymentService(db fixtureDB) (*deployment.Service, error) { - rules, err := placement.NewRules(providers.Builtin(), db.publicURL) - if err != nil { - return nil, err - } - adapter := deploymentpg.New(pgunit.NewPool(db.pool), db.cipher) - return deployment.NewService(adapter, adapter, providers.Builtin(), rules) -} - -// fixtureRules builds the placement rules on db's public URL, as cmd/server -// does for the Store's Session creation. -func fixtureRules(t testing.TB, db fixtureDB) *placement.Rules { - t.Helper() - rules, err := placement.NewRules(providers.Builtin(), db.publicURL) - if err != nil { - t.Fatal(err) - } - return rules -} - -// fixtureReader is the deployment reader on db, as cmd/server builds it. -func fixtureReader(db fixtureDB) *deploymentpg.Store { - return deploymentpg.New(pgunit.NewPool(db.pool), db.cipher) +// fixtureDeploymentService builds the deployment service as cmd/server does, on +// s's database, credential key and placement rules. +func fixtureDeploymentService(s *Store) (*deployment.Service, error) { + adapter := deploymentStore(s) + return deployment.NewService(adapter, adapter, providers.Builtin(), s.placement) } // fixtureOwnerEpoch reads the execution owner epoch from the deployment store, // as cmd/server does to fence node connections. -func fixtureOwnerEpoch(t testing.TB, db fixtureDB) uint64 { +func fixtureOwnerEpoch(t testing.TB, s *Store) uint64 { t.Helper() - epoch, err := deploymentpg.New(pgunit.NewPool(db.pool), db.cipher).OwnerEpoch(t.Context()) + epoch, err := deploymentStore(s).OwnerEpoch(t.Context()) if err != nil { t.Fatal(err) } @@ -58,11 +29,11 @@ func fixtureOwnerEpoch(t testing.TB, db fixtureDB) uint64 { // fixtureDeploymentExecution builds the pooled deployment service and the // deployment execution operations on lease, as cmd/server does for the Worker. -func fixtureDeploymentExecution(db fixtureDB, lease *pgunit.Lease) (*deployment.Service, *deployment.ExecutionOperations, error) { - service, err := fixtureDeploymentService(db) +func fixtureDeploymentExecution(s *Store, lease *pgunit.Lease) (*deployment.Service, *deployment.ExecutionOperations, error) { + service, err := fixtureDeploymentService(s) if err != nil { return nil, nil, err } - changes, err := deployment.NewExecutionOperations(service, deploymentpg.NewExecution(lease, db.cipher)) + changes, err := deployment.NewExecutionOperations(service, deploymentpg.NewExecution(lease, s.credentialCipher)) return service, changes, err } diff --git a/services/core/tests/integration/dispatch_test.go b/services/core/tests/integration/dispatch_test.go index e4b122c3b..40de3c693 100644 --- a/services/core/tests/integration/dispatch_test.go +++ b/services/core/tests/integration/dispatch_test.go @@ -30,7 +30,6 @@ type dispatchHarness struct { admissions map[string]fixtureAdmission t *testing.T s *Store - db fixtureDB d *execution.Dispatcher lease execution.Ownership // held by tests that run execution operations without a Worker owned *execution.Owner // the Owner that bound binds, acquired on first use @@ -51,8 +50,8 @@ func newDispatchHarness(t *testing.T) *dispatchHarness { func newDispatchHarnessForSession(t *testing.T, configuration []byte, local bool) *dispatchHarness { t.Helper() - s, db := newModelTestStoreDB(t) - h := &dispatchHarness{t: t, s: s, db: db, tenant: uuid.NewString(), environments: map[string]*dispatchHarness{}} + s, _ := NewModelTestStore(t) + h := &dispatchHarness{t: t, s: s, tenant: uuid.NewString(), environments: map[string]*dispatchHarness{}} ctx := context.Background() var err error h.session, err = s.CreateSession(ctx, h.tenant, WithFixtureModelProvider(sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "session", Configuration: configuration})) @@ -68,26 +67,26 @@ func newDispatchHarnessForSession(t *testing.T, configuration []byte, local bool } _ = json.Unmarshal(configuration, &snapshot) if snapshot.Environment.Type == "self_hosted" { - h.device, h.credential = enrollFixtureSession(t, s, db, h.tenant, h.session) + h.device, h.credential = enrollFixtureSession(t, s, h.tenant, h.session) secret = h.credential } else if local { - environment, getErr := fixtureSessionStore(db).GetSessionEnvironment(ctx, h.tenant, h.session.ID) + environment, getErr := sessionAdapter(s).GetSessionEnvironment(ctx, h.tenant, h.session.ID) if getErr != nil { t.Fatal(getErr) } - h.device, err = FixtureEnvironmentDevice(ctx, db.pool, h.tenant, environment.ID, "local runtime", runtimedevice.HashCredential(secret)) + h.device, err = FixtureEnvironmentDevice(ctx, s.pool, h.tenant, environment.ID, "local runtime", runtimedevice.HashCredential(secret)) } else { - h.device, err = fixtureSessionService(t, db).CreateDevice(ctx, h.tenant, "isolated executor", runtimedevice.HashCredential(secret)) + h.device, err = sessionService(t, s).CreateDevice(ctx, h.tenant, "isolated executor", runtimedevice.HashCredential(secret)) } if err != nil { t.Fatal(err) } - if err = bindSessionDevice(t, db, h.tenant, h.session.ID, h.device.ID); err != nil { + if err = bindSessionDevice(t, s, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - server.Config.Handler, h.registry, err = runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), fixtureSessionStore(db), wsURL) + server.Config.Handler, h.registry, err = runtime.NewGateway(sessionAdapter(s), sessionService(t, s), sessionAdapter(s), wsURL) if err != nil { t.Fatal(err) } @@ -115,11 +114,12 @@ func newDispatchHarnessForSession(t *testing.T, configuration []byte, local bool } time.Sleep(10 * time.Millisecond) } - sessionStore, sessionService, err := fixtureSessions(db) + sessionStore := sessionAdapter(s) + sessionService, err := newSessionService(s) if err != nil { t.Fatal(err) } - h.d = &execution.Dispatcher{Registry: h.registry, Observer: modelconfigurationpg.New(pgunit.NewPool(db.pool), db.cipher), Sessions: sessionService, SessionsReader: sessionStore} + h.d = &execution.Dispatcher{Registry: h.registry, Observer: modelconfigurationpg.New(pgunit.NewPool(s.pool), s.credentialCipher), Sessions: sessionService, SessionsReader: sessionStore} return h } @@ -182,7 +182,7 @@ type runResult struct { func (h *dispatchHarness) owner() execution.Owner { h.t.Helper() if h.owned == nil { - owner := executionOwner(h.t, h.db) + owner := executionOwner(h.t, h.s) h.owned = &owner } return *h.owned @@ -259,13 +259,13 @@ func TestExecutionDispatchSteeringAndNativeContinuity(t *testing.T) { t.Fatalf("missing result: %+v", outcome) } // Restart Core: a new Store and execution owner continue the native Session. - awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, h.db.pool) + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, h.s.pool) if err := h.owner().Lease.Close(ctx); err != nil { t.Fatal(err) } awaitRelease() - newStore, db := newTestStoreDB(t) - h.s, h.db, h.owned = newStore, db, nil + newStore, _ := testStore(t) + h.s, h.owned = newStore, nil bound, err := sessionAdapter(newStore).GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || bound.NativeSessionID != "native-thread-1" { t.Fatalf("native binding lost: %+v %v", bound, err) diff --git a/services/core/tests/integration/environment_admission_test.go b/services/core/tests/integration/environment_admission_test.go index 8229152ba..27e191352 100644 --- a/services/core/tests/integration/environment_admission_test.go +++ b/services/core/tests/integration/environment_admission_test.go @@ -23,7 +23,7 @@ func newEnvironmentAdmission(t *testing.T) (*dispatchHarness, *execution.Worker) t.Helper() h := newDispatchHarness(t) enableWorkerEnvironment(t, h) - worker := startWorker(t, t.Context(), h.db, h.d) + worker := startWorker(t, t.Context(), h.s, h.d) t.Cleanup(func() { ctx, cancel := context.WithCancel(context.Background()) cancel() diff --git a/services/core/tests/integration/environment_claim_worker_test.go b/services/core/tests/integration/environment_claim_worker_test.go index e7fd2c54e..522f5230a 100644 --- a/services/core/tests/integration/environment_claim_worker_test.go +++ b/services/core/tests/integration/environment_claim_worker_test.go @@ -15,9 +15,9 @@ import ( func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { for _, deleted := range []bool{false, true} { t.Run(map[bool]string{false: "unbound", true: "deleted"}[deleted], func(t *testing.T) { - s, db := newTestStoreDB(t) + s, _ := testStore(t) tenant, pending := newEnvironmentExpiryReservation(t, s) - owner := executionOwner(t, db) + owner := executionOwner(t, s) got, err := owner.Sessions.PromoteEnvironmentInput(t.Context(), tenant, pending.SessionID, pending.ID) if err != nil || len(got.Receipts) != 1 || got.Receipts[0].Replayed { t.Fatal(got, err) @@ -36,15 +36,15 @@ func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { t.Fatal("promotion did not retain the active claim", turn, err) } // Simulate owner loss after commit, without sending any daemon Start. - awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, db.pool) + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, s.pool) if err := owner.Lease.Close(t.Context()); err != nil { t.Fatal(err) } awaitRelease() - restarted := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry()}) + restarted := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry()}) stopped, cancel := context.WithCancel(t.Context()) cancel() - awaitRelease = pgtest.ObserveExecutionLeaseRelease(t, db.pool) + awaitRelease = pgtest.ObserveExecutionLeaseRelease(t, s.pool) if err := restarted.Run(stopped); !errors.Is(err, context.Canceled) { t.Fatal(err) } @@ -57,14 +57,14 @@ func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { t.Fatal("restart failed to settle the original claim", turn, err) } var turns, inputs, queued int - err = db.pool.QueryRow(t.Context(), `SELECT + err = s.pool.QueryRow(t.Context(), `SELECT (SELECT count(*) FROM turns WHERE session_id=$1), (SELECT count(*) FROM turn_inputs WHERE session_id=$1), (SELECT count(*) FROM turns WHERE session_id=$1 AND status='queued')`, pending.SessionID).Scan(&turns, &inputs, &queued) if err != nil || turns != 1 || inputs != 1 || queued != 0 { t.Fatal("restart duplicated or requeued prepared work", turns, inputs, queued, err) } - successor := executionOwner(t, db).Sessions + successor := executionOwner(t, s).Sessions retry, err := successor.PromoteEnvironmentInput(t.Context(), tenant, pending.SessionID, pending.ID) if deleted { if !errors.Is(err, sessions.ErrNotFound) { diff --git a/services/core/tests/integration/environment_connection_events_test.go b/services/core/tests/integration/environment_connection_events_test.go index 7096f339e..a3bfa945f 100644 --- a/services/core/tests/integration/environment_connection_events_test.go +++ b/services/core/tests/integration/environment_connection_events_test.go @@ -6,14 +6,13 @@ import ( "time" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" - "github.com/jackc/pgx/v5/pgxpool" ) -func awaitEnvironmentConnectionState(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tenant, environment, status string) { +func awaitEnvironmentConnectionState(t *testing.T, ctx context.Context, s *Store, tenant, environment, status string) { t.Helper() deadline := time.Now().Add(5 * time.Second) for time.Now().Before(deadline) { - value, err := sessionReads(pool).GetEnvironment(ctx, tenant, environment) + value, err := sessionAdapter(s).GetEnvironment(ctx, tenant, environment) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/environment_connection_worker_test.go b/services/core/tests/integration/environment_connection_worker_test.go index a5d855f51..ca356f402 100644 --- a/services/core/tests/integration/environment_connection_worker_test.go +++ b/services/core/tests/integration/environment_connection_worker_test.go @@ -14,17 +14,17 @@ import ( ) func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { - s, db := newTestStoreDB(t) + s, _ := testStore(t) tenant := uuid.NewString() session, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "connection-worker", Configuration: []byte(`{"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`)}) if err != nil { t.Fatal(err) } - environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } - owner := executionOwner(t, db) + owner := executionOwner(t, s) generation := uuid.NewString() if err := owner.Sessions.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { t.Fatal(err) @@ -32,12 +32,12 @@ func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { if err := owner.Sessions.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err != nil { t.Fatal(err) } - awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, db.pool) + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, s.pool) if err := owner.Lease.Close(t.Context()); err != nil { t.Fatal(err) } awaitRelease() - worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry()}) + worker := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry()}) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) exited := make(chan struct{}) @@ -50,7 +50,7 @@ func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { t.Error("worker cleanup did not exit") } }) - awaitEnvironmentConnectionState(t, ctx, db.pool, tenant, environment.ID, "disconnected") + awaitEnvironmentConnectionState(t, ctx, s, tenant, environment.ID, "disconnected") cancel() select { @@ -61,7 +61,7 @@ func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("worker did not close") } - awaitEnvironmentConnectionState(t, t.Context(), db.pool, tenant, environment.ID, "disconnected") + awaitEnvironmentConnectionState(t, t.Context(), s, tenant, environment.ID, "disconnected") if err := worker.CheckOwnership(t.Context()); err == nil { t.Fatal("worker retained lease") } diff --git a/services/core/tests/integration/environment_device_test.go b/services/core/tests/integration/environment_device_test.go index 493be5031..b348101ea 100644 --- a/services/core/tests/integration/environment_device_test.go +++ b/services/core/tests/integration/environment_device_test.go @@ -25,7 +25,7 @@ func TestWorkerEnvironmentSelectsCapableDeviceWithoutMovingBinding(t *testing.T) awaitFixtureCapabilities(t, originalRuntime, caps) generalFrames := workerFrames(t, h) boundFrames := workerFrames(t, originalRuntime) - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) select { case frame := <-generalFrames: t.Fatal("general device received self-hosted work", frame.Type) @@ -33,7 +33,7 @@ func TestWorkerEnvironmentSelectsCapableDeviceWithoutMovingBinding(t *testing.T) t.Fatal("incapable enrolled device received work", frame.Type) case <-time.After(time.Second): } - if _, err := fixtureSessionStore(h.db).GetSessionDevice(t.Context(), h.tenant, pending.SessionID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(h.s).GetSessionDevice(t.Context(), h.tenant, pending.SessionID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unregistered Runtime was assigned general compute", err) } session, err := sessionAdapter(h.s).GetSession(t.Context(), h.tenant, pending.SessionID) @@ -43,11 +43,11 @@ func TestWorkerEnvironmentSelectsCapableDeviceWithoutMovingBinding(t *testing.T) other := connectFixtureRuntime(t, h, session) otherFrames := workerFrames(t, other) request := nextWorkerFrame(t, otherFrames, proto.TypeExecutionPrepare) - selected, err := fixtureSessionStore(h.db).GetSessionDevice(t.Context(), h.tenant, pending.SessionID) + selected, err := sessionAdapter(h.s).GetSessionDevice(t.Context(), h.tenant, pending.SessionID) if err != nil || selected.ID != other.device.ID { t.Fatal("enrollment did not retain exact Runtime", err) } - original, err := fixtureSessionStore(h.db).GetSessionDevice(t.Context(), h.tenant, bound.SessionID) + original, err := sessionAdapter(h.s).GetSessionDevice(t.Context(), h.tenant, bound.SessionID) if err != nil || original.ID != originalRuntime.device.ID { t.Fatal("existing binding moved to a capable Runtime", err) } diff --git a/services/core/tests/integration/environment_directory_test.go b/services/core/tests/integration/environment_directory_test.go index aebfd8f0b..815203094 100644 --- a/services/core/tests/integration/environment_directory_test.go +++ b/services/core/tests/integration/environment_directory_test.go @@ -21,7 +21,7 @@ type directoryResult struct { func directoryWorker(t *testing.T) (*dispatchHarness, *execution.Worker, sessions.Environment) { t.Helper() h := newDispatchHarnessForSession(t, []byte(`{"agent":{"model":"unavailable-model"},"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`), true) - environment, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) + environment, err := sessionAdapter(h.s).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) if err != nil { t.Fatal(err) } @@ -34,7 +34,7 @@ func directoryWorker(t *testing.T) (*dispatchHarness, *execution.Worker, session info, _, _ := peer.AgentKindStatus("codex") return info.Capabilities.WorkspaceReadPreparation }) - w := startWorker(t, t.Context(), h.db, h.d) + w := startWorker(t, t.Context(), h.s, h.d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- w.Run(ctx) }() diff --git a/services/core/tests/integration/environment_executor_command_test.go b/services/core/tests/integration/environment_executor_command_test.go index c0a906dcb..2803532f5 100644 --- a/services/core/tests/integration/environment_executor_command_test.go +++ b/services/core/tests/integration/environment_executor_command_test.go @@ -52,7 +52,7 @@ func TestEnvironmentExecutorOperatorCommand(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := sessionReads(pool).GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -64,14 +64,14 @@ func TestEnvironmentExecutorOperatorCommand(t *testing.T) { if next == first { t.Fatal("rotation returned the same key") } - if _, err := sessionReads(pool).AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(first)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(first)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("old command credential retained authority", err) } - if owner, err := sessionReads(pool).AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(next)); err != nil || owner != tenant { + if owner, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(next)); err != nil || owner != tenant { t.Fatal("rotated command credential failed", err) } command(tenant, true, "--revoke") - if _, err := sessionReads(pool).AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(next)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(next)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("revoked command credential retained authority", err) } t.Log("built operator command issued before Session creation, rejected duplicate/foreign requests, rotated and revoked durable credentials") diff --git a/services/core/tests/integration/environment_expiry_dispatch_test.go b/services/core/tests/integration/environment_expiry_dispatch_test.go index 7b57eb2be..623fe0ea5 100644 --- a/services/core/tests/integration/environment_expiry_dispatch_test.go +++ b/services/core/tests/integration/environment_expiry_dispatch_test.go @@ -19,7 +19,7 @@ func TestWorkerEnvironmentExpiryAtFullExecutionCapacity(t *testing.T) { h := newDispatchHarness(t) _, pool := testStore(t) enableEnvironmentExpiryDispatch(h) - worker, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + worker, stop := startEnvironmentExpiryWorker(t, h.s, h.d) var requests []proto.Envelope var active []sessions.Session for _, key := range []string{"one", "two", "three", "four"} { @@ -64,7 +64,7 @@ func TestWorkerEnvironmentExpirySkipsBusySessionAndAllowsDispatch(t *testing.T) if _, err := tx.Exec(t.Context(), "SELECT id FROM sessions WHERE id=$1 FOR UPDATE", locked.SessionID); err != nil { t.Fatal(err) } - worker, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + worker, stop := startEnvironmentExpiryWorker(t, h.s, h.d) h.session = publicSession(t, h, "unrelated") receipt, err := worker.SubmitInputs(t.Context(), h.tenant, h.session.ID, "work", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"make normal progress"}`)}}) if err != nil { diff --git a/services/core/tests/integration/environment_expiry_worker_test.go b/services/core/tests/integration/environment_expiry_worker_test.go index 44892e17d..a7d8b4fdd 100644 --- a/services/core/tests/integration/environment_expiry_worker_test.go +++ b/services/core/tests/integration/environment_expiry_worker_test.go @@ -39,9 +39,9 @@ func makeEnvironmentExpiryDue(t *testing.T, pool *pgxpool.Pool, pending *session } } -func startEnvironmentExpiryWorker(t *testing.T, db fixtureDB, d *execution.Dispatcher) (*execution.Worker, func()) { +func startEnvironmentExpiryWorker(t *testing.T, s *Store, d *execution.Dispatcher) (*execution.Worker, func()) { t.Helper() - worker := startWorker(t, t.Context(), db, d) + worker := startWorker(t, t.Context(), s, d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -99,22 +99,22 @@ func assertEnvironmentExpiryHasNoHistory(t *testing.T, pool *pgxpool.Pool, sessi } func TestWorkerEnvironmentExpiryWithoutDevicesAndAfterRestart(t *testing.T) { - s, db := newTestStoreDB(t) + s, _ := testStore(t) dueTenant, due := newEnvironmentExpiryReservation(t, s) futureTenant, future := newEnvironmentExpiryReservation(t, s) - makeEnvironmentExpiryDue(t, db.pool, &due) + makeEnvironmentExpiryDue(t, s.pool, &due) d := &execution.Dispatcher{Registry: runtimegateway.NewRegistry()} - _, stop := startEnvironmentExpiryWorker(t, db, d) + _, stop := startEnvironmentExpiryWorker(t, s, d) waitEnvironmentExpiry(t, s, dueTenant, due) got, err := sessionAdapter(s).GetEnvironmentInputReservation(t.Context(), futureTenant, future.SessionID, future.ID) if err != nil || got.State != sessions.EnvironmentInputPending || !got.Deadline.Equal(future.Deadline) { t.Fatal("future input changed", got, err) } stop() - makeEnvironmentExpiryDue(t, db.pool, &future) - _, stop = startEnvironmentExpiryWorker(t, db, d) + makeEnvironmentExpiryDue(t, s.pool, &future) + _, stop = startEnvironmentExpiryWorker(t, s, d) waitEnvironmentExpiry(t, s, futureTenant, future) stop() - assertEnvironmentExpiryHasNoHistory(t, db.pool, due.SessionID) - assertEnvironmentExpiryHasNoHistory(t, db.pool, future.SessionID) + assertEnvironmentExpiryHasNoHistory(t, s.pool, due.SessionID) + assertEnvironmentExpiryHasNoHistory(t, s.pool, future.SessionID) } diff --git a/services/core/tests/integration/environment_file_write_semantics_public_test.go b/services/core/tests/integration/environment_file_write_semantics_public_test.go index 949e7206c..03398dc75 100644 --- a/services/core/tests/integration/environment_file_write_semantics_public_test.go +++ b/services/core/tests/integration/environment_file_write_semantics_public_test.go @@ -59,13 +59,13 @@ func TestEnvironmentFileCreateRejectionsLeaveNoReceiptOrConsumption(t *testing.T {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "tenant-b", TokenSHA256: runtimedevice.HashCredential(other), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(t, w)) + handler, err := publicHandler(t, h.s, auth, "codex", workerExecution(t, w)) if err != nil { t.Fatal(err) } server := httptest.NewServer(handler) defer server.Close() - fileStore, fileService := fixtureFiles(t, h.db) + fileStore, fileService := fixtureFiles(t, h.s) source, err := fileService.Create(t.Context(), files.CreateCommand{TenantID: h.tenant, Upload: func(out io.Writer) (files.Upload, error) { _, err := out.Write([]byte("src")) return files.Upload{Filename: "source.txt", Purpose: files.PurposeUserData}, err @@ -140,7 +140,7 @@ func TestEnvironmentFileCreateRejectionsLeaveNoReceiptOrConsumption(t *testing.T done := post(token, environment.ID, tc.body) id := serveFileWrite(h, rejected(tc.reason)) assertError(await(done), tc.message) - if intent, err := FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, id); err != nil || intent.State != "rejected" { + if intent, err := FixtureFileWrite(t.Context(), h.s.pool, h.tenant, environment.ID, id); err != nil || intent.State != "rejected" { t.Fatal("rejection did not settle", intent, err) } } diff --git a/services/core/tests/integration/environment_initial_public_test.go b/services/core/tests/integration/environment_initial_public_test.go index 12c062565..a3b47fa76 100644 --- a/services/core/tests/integration/environment_initial_public_test.go +++ b/services/core/tests/integration/environment_initial_public_test.go @@ -21,7 +21,7 @@ func TestEnvironmentInitialFailureOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, db := newTestStoreDB(t) + s, _ := testStore(t) tenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, @@ -36,7 +36,7 @@ func TestEnvironmentInitialFailureOfficialClient(t *testing.T) { if err != nil { t.Fatal(err) } - owner := executionOwner(t, db) + owner := executionOwner(t, s) t.Cleanup(func() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() @@ -44,7 +44,7 @@ func TestEnvironmentInitialFailureOfficialClient(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) + handler, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } @@ -89,7 +89,7 @@ func TestEnvironmentInitialFailureOfficialClient(t *testing.T) { } } var reservation string - if err := db.pool.QueryRow(t.Context(), "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1 AND is_initial RETURNING id", session.ID).Scan(&reservation); err != nil { + if err := s.pool.QueryRow(t.Context(), "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1 AND is_initial RETURNING id", session.ID).Scan(&reservation); err != nil { t.Fatal(err) } if result, err := sessionService(t, s).ExpireEnvironmentInput(t.Context(), tenant, session.ID, reservation); err != nil || result.State != sessions.EnvironmentInputExpired { diff --git a/services/core/tests/integration/environment_initialization_test.go b/services/core/tests/integration/environment_initialization_test.go index a18d50f6d..d21de7390 100644 --- a/services/core/tests/integration/environment_initialization_test.go +++ b/services/core/tests/integration/environment_initialization_test.go @@ -20,26 +20,25 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/google/uuid" - "github.com/jackc/pgx/v5/pgxpool" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -func initializationState(t *testing.T, pool *pgxpool.Pool, tenant, environment string) string { +func initializationState(t *testing.T, s *Store, tenant, environment string) string { t.Helper() - value, err := sessionReads(pool).GetEnvironment(t.Context(), tenant, environment) + value, err := sessionAdapter(s).GetEnvironment(t.Context(), tenant, environment) if err != nil { t.Fatal(err) } return value.Initialization } -func awaitInitialization(t *testing.T, pool *pgxpool.Pool, tenant, environment, state string) { +func awaitInitialization(t *testing.T, s *Store, tenant, environment, state string) { t.Helper() deadline := time.Now().Add(10 * time.Second) for time.Now().Before(deadline) { - if initializationState(t, pool, tenant, environment) == state { + if initializationState(t, s, tenant, environment) == state { return } time.Sleep(10 * time.Millisecond) @@ -55,7 +54,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) principal := FixtureExecutorPrincipal(t, s, uuid.NewString()) session, err := s.CreateSession(t.Context(), principal.TenantID, sessions.CreateSession{ Creator: principal.Subject(), Engine: "codex", IdempotencyKey: uuid.NewString(), @@ -66,23 +65,23 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test if err != nil { t.Fatal(err) } - environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) if err != nil { t.Fatal(err) } - key, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + key, err := sessionService(t, s).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } - enrolled, err := fixtureSessionService(t, db).EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) + enrolled, err := sessionService(t, s).EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) if err != nil { t.Fatal(err) } registry := runtimegateway.NewRegistry() - handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(fixtureSessionStore(db)), Registry: registry}) + handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(sessionAdapter(s)), Registry: registry}) server := httptest.NewServer(http.HandlerFunc(handler.WS)) defer server.Close() - worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: registry}) + worker := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: registry}) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -111,7 +110,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test actions = append(actions, action) mu.Unlock() if outcome == "revoked" { - if err := fixtureSessionService(t, db).RevokeDevice(t.Context(), principal.TenantID, enrolled.DeviceID); err != nil { + if err := sessionService(t, s).RevokeDevice(t.Context(), principal.TenantID, enrolled.DeviceID); err != nil { t.Error(err) } return completedInitialization(request, data) @@ -122,7 +121,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test return completedInitialization(request, data) } time.Sleep(350 * time.Millisecond) - if initializationState(t, db.pool, principal.TenantID, environment.ID) != "pending" { + if initializationState(t, s, principal.TenantID, environment.ID) != "pending" { t.Fatal("unconnected preparation was consumed") } bootstrap := sandbox.Bootstrap{DeviceID: enrolled.DeviceID, Credential: key.Token} @@ -133,8 +132,8 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test if outcome != "completed" { want = "failed" } - awaitInitialization(t, db.pool, principal.TenantID, environment.ID, want) - if _, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: principal.TenantID, EnvironmentID: environment.ID}); !errors.Is(err, deployment.ErrNotFound) { + awaitInitialization(t, s, principal.TenantID, environment.ID, want) + if _, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: principal.TenantID, EnvironmentID: environment.ID}); !errors.Is(err, deployment.ErrNotFound) { t.Fatal("self-hosted preparation fabricated allocation", err) } if outcome == "completed" { @@ -171,7 +170,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test // Revocation can commit after the worker observes a connected peer but before // it claims preparation. It must remain a per-Environment admission result. func TestEnvironmentInitializationRevocationBeforeClaim(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) principal := FixtureExecutorPrincipal(t, s, uuid.NewString()) create := func() sessions.EnvironmentInitialization { t.Helper() @@ -183,23 +182,23 @@ func TestEnvironmentInitializationRevocationBeforeClaim(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) if err != nil { t.Fatal(err) } - key, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + key, err := sessionService(t, s).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } - enrolled, err := fixtureSessionService(t, db).EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) + enrolled, err := sessionService(t, s).EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) if err != nil { t.Fatal(err) } return sessions.EnvironmentInitialization{EnvironmentID: environment.ID, SessionID: session.ID, TenantID: principal.TenantID, DeviceID: enrolled.DeviceID, State: "pending", Engine: "codex"} } revoked, other := create(), create() - owned := executionOwner(t, db).Sessions - if err := fixtureSessionService(t, db).RevokeDevice(t.Context(), principal.TenantID, revoked.DeviceID); err != nil { + owned := executionOwner(t, s).Sessions + if err := sessionService(t, s).RevokeDevice(t.Context(), principal.TenantID, revoked.DeviceID); err != nil { t.Fatal(err) } if err := owned.ClaimEnvironmentInitialization(t.Context(), revoked); !errors.Is(err, sessions.ErrNotFound) { @@ -210,7 +209,7 @@ func TestEnvironmentInitializationRevocationBeforeClaim(t *testing.T) { if err := owned.ClaimEnvironmentInitialization(t.Context(), stale); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatalf("stale binding escaped normal admission handling: %v", err) } - if initializationState(t, db.pool, other.TenantID, other.EnvironmentID) != "pending" { + if initializationState(t, s, other.TenantID, other.EnvironmentID) != "pending" { t.Fatal("stale claim changed preparation state") } if err := owned.ClaimEnvironmentInitialization(t.Context(), other); err != nil { diff --git a/services/core/tests/integration/environment_mcp_public_test.go b/services/core/tests/integration/environment_mcp_public_test.go index e616943e9..d57a013fd 100644 --- a/services/core/tests/integration/environment_mcp_public_test.go +++ b/services/core/tests/integration/environment_mcp_public_test.go @@ -14,9 +14,9 @@ import ( func TestPublicEnvironmentMCPUsesAttachedVaultSelection(t *testing.T) { for _, kind := range []string{"codex", "claude_sdk", "mcode"} { t.Run(kind, func(t *testing.T) { - s, db, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) + s, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "test", TenantID: tenant, TokenSHA256: runtimedevice.HashCredential("test-token")}}) - handler, err := publicHandler(t, s, db, auth, kind, workerExecution(t, &execution.Worker{}), executorURL("https://executor.example")) + handler, err := publicHandler(t, s, auth, kind, workerExecution(t, &execution.Worker{}), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } @@ -47,7 +47,7 @@ func TestPublicEnvironmentMCPUsesAttachedVaultSelection(t *testing.T) { } } } - assertSelfHostedMCPRejectionHasNoWrites(t, db.pool, tenant) + assertSelfHostedMCPRejectionHasNoWrites(t, s.pool, tenant) for _, r := range []*httptest.ResponseRecorder{send("environment", []string{}, nil, nil, false), send("environment", []string{vault.ID}, credential.ID, nil, false), send("environment", []string{vault.ID}, nil, nil, false)} { if r.Code != 201 { t.Fatal("qualified public MCP rejected", r.Code, r.Body) diff --git a/services/core/tests/integration/environment_plugins_test.go b/services/core/tests/integration/environment_plugins_test.go index 030ffec5e..21d207470 100644 --- a/services/core/tests/integration/environment_plugins_test.go +++ b/services/core/tests/integration/environment_plugins_test.go @@ -21,7 +21,7 @@ func TestPluginsFrozenInSession(t *testing.T) { if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) var archive bytes.Buffer writer := zip.NewWriter(&archive) for path, body := range map[string]string{ diff --git a/services/core/tests/integration/environment_retrieve_public_test.go b/services/core/tests/integration/environment_retrieve_public_test.go index 6223410f6..2169efcef 100644 --- a/services/core/tests/integration/environment_retrieve_public_test.go +++ b/services/core/tests/integration/environment_retrieve_public_test.go @@ -21,7 +21,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, db := newModelTestStoreDB(t) + s, _ := NewModelTestStore(t) tenant, foreignTenant := uuid.NewString(), uuid.NewString() principal := FixtureExecutorPrincipal(t, s, tenant) token, peer, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() @@ -33,7 +33,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { if err := s.EnsureProjectScopes(t.Context(), []identity.ProjectScope{{TenantID: tenant, OrganizationID: principal.OrganizationID, ProjectID: tenant}, {TenantID: foreignTenant, OrganizationID: principal.OrganizationID, ProjectID: foreignTenant}}); err != nil { t.Fatal(err) } - executor, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), "") + executor, err := sessionService(t, s).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), "") if err != nil { t.Fatal(err) } @@ -42,12 +42,12 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { if !revoked { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if err := fixtureSessionService(t, db).RevokeExecutorCredential(ctx, principal, executor.KeyID); err != nil { + if err := sessionService(t, s).RevokeExecutorCredential(ctx, principal, executor.KeyID); err != nil { t.Error("owned executor credential cleanup failed", err) } } }() - handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://private-registry.example")) + handler, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), executorURL("https://private-registry.example")) if err != nil { t.Fatal(err) } @@ -75,18 +75,18 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { return result } result := run() - before, err := fixtureSessionStore(db).GetEnvironment(t.Context(), tenant, result["environment_id"]) + before, err := sessionAdapter(s).GetEnvironment(t.Context(), tenant, result["environment_id"]) if err != nil { t.Fatal(err) } - if err := fixtureSessionService(t, db).RevokeExecutorCredential(t.Context(), principal, executor.KeyID); err != nil { + if err := sessionService(t, s).RevokeExecutorCredential(t.Context(), principal, executor.KeyID); err != nil { t.Fatal(err) } revoked = true server.Close() - db.pool.Close() - reopened, reopenedDB := newModelTestStoreDB(t) - handler, err = publicHandler(t, reopened, reopenedDB, auth, "codex") + s.pool.Close() + reopened, _ := NewModelTestStore(t) + handler, err = publicHandler(t, reopened, auth, "codex") if err != nil { t.Fatal(err) } @@ -97,12 +97,12 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { settings[key] = value } run() - after, err := fixtureSessionStore(reopenedDB).GetEnvironment(t.Context(), tenant, before.ID) + after, err := sessionAdapter(reopened).GetEnvironment(t.Context(), tenant, before.ID) if err != nil || !reflect.DeepEqual(before, after) { t.Fatal("public retrieval changed durable Environment state", err) } var history int - if err := reopenedDB.pool.QueryRow(t.Context(), `SELECT + if err := reopened.pool.QueryRow(t.Context(), `SELECT (SELECT count(*) FROM turns WHERE session_id=$1) + (SELECT count(*) FROM environment_input_reservations WHERE session_id=$1) + (SELECT count(*) FROM session_events WHERE session_id=$1)`, before.SessionID).Scan(&history); err != nil || history != 0 { diff --git a/services/core/tests/integration/environment_runtime_fixture_test.go b/services/core/tests/integration/environment_runtime_fixture_test.go index 1b35f5d11..4a047b13a 100644 --- a/services/core/tests/integration/environment_runtime_fixture_test.go +++ b/services/core/tests/integration/environment_runtime_fixture_test.go @@ -13,22 +13,22 @@ import ( "github.com/gorilla/websocket" ) -func enrollFixtureSession(t *testing.T, s *Store, db fixtureDB, tenant string, session sessions.Session) (sessions.ExecutionDevice, string) { +func enrollFixtureSession(t *testing.T, s *Store, tenant string, session sessions.Session) (sessions.ExecutionDevice, string) { t.Helper() - environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } principal := FixtureExecutorPrincipal(t, s, tenant) - key, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + key, err := sessionService(t, s).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } - enrolled, err := fixtureSessionService(t, db).EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) + enrolled, err := sessionService(t, s).EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) if err != nil || enrolled.EnvironmentID != environment.ID || enrolled.SessionID != session.ID || enrolled.WorkspaceDirectory != "/workspace" { t.Fatalf("Runtime enrollment: %+v %v", enrolled, err) } - bound, err := fixtureSessionStore(db).GetSessionDevice(t.Context(), tenant, session.ID) + bound, err := sessionAdapter(s).GetSessionDevice(t.Context(), tenant, session.ID) if err != nil || bound.ID != enrolled.DeviceID || bound.EnvironmentID != environment.ID { t.Fatalf("Runtime binding: %+v %v", bound, err) } @@ -38,9 +38,9 @@ func enrollFixtureSession(t *testing.T, s *Store, db fixtureDB, tenant string, s func connectFixtureRuntime(t *testing.T, h *dispatchHarness, session sessions.Session) *dispatchHarness { t.Helper() // The Runtime shares the harness's Core, not its connection or write lock. - other := &dispatchHarness{t: h.t, s: h.s, db: h.db, lease: h.lease, owned: h.owned, d: h.d, tenant: h.tenant, session: session, registry: h.registry, url: h.url, + other := &dispatchHarness{t: h.t, s: h.s, lease: h.lease, owned: h.owned, d: h.d, tenant: h.tenant, session: session, registry: h.registry, url: h.url, admissions: h.admissions, environments: h.environments} - other.device, other.credential = enrollFixtureSession(t, h.s, h.db, h.tenant, session) + other.device, other.credential = enrollFixtureSession(t, h.s, h.tenant, session) u, err := url.Parse(h.url) if err != nil { t.Fatal(err) @@ -74,7 +74,7 @@ func completeEmptyArtifactExport(t *testing.T, h *dispatchHarness, frames ...<-c } frame := read(proto.TypeExecutionPrepare) var prepare proto.ExecutionPreparePayload - environment, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) + environment, err := sessionAdapter(h.s).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) if err != nil || frame.DecodePayload(&prepare) != nil || !proto.ValidWorkspaceReadPreparation(prepare.Configuration) || prepare.Configuration.LocalEnvironment == nil || prepare.Configuration.LocalEnvironment.ID != environment.ID { t.Fatal("artifact preparation lost exact local authority", err) } diff --git a/services/core/tests/integration/environment_setup_test.go b/services/core/tests/integration/environment_setup_test.go index c033371f3..039131e0e 100644 --- a/services/core/tests/integration/environment_setup_test.go +++ b/services/core/tests/integration/environment_setup_test.go @@ -20,7 +20,7 @@ func TestEnvironmentSetupEncryptedSnapshotAndIsolation(t *testing.T) { if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) tenant, foreign := uuid.NewString(), uuid.NewString() setup := environmentconfig.Setup{Env: map[string]string{"SECRET": "session-env-canary"}, Commands: []environmentconfig.SetupCommand{{Command: "printf session-command-canary > result"}}, Packages: v1.EnvironmentPackages{NPM: []string{"is-number@7.0.0"}}} input := sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"environment":{"type":"openai_hosted"}}`), Initialization: setup} diff --git a/services/core/tests/integration/environment_work_test.go b/services/core/tests/integration/environment_work_test.go index d8efad184..fae8b6abc 100644 --- a/services/core/tests/integration/environment_work_test.go +++ b/services/core/tests/integration/environment_work_test.go @@ -40,7 +40,7 @@ func TestEnvironmentInputWorkFiltersAndPagesDevices(t *testing.T) { default: runtime := h.environments[pending.SessionID] if state == "revoked" { - if err := fixtureSessionService(t, h.db).RevokeDevice(t.Context(), h.tenant, runtime.device.ID); err != nil { + if err := sessionService(t, h.s).RevokeDevice(t.Context(), h.tenant, runtime.device.ID); err != nil { t.Fatal(err) } work, err := sessionAdapter(h.s).ListEnvironmentInputWork(t.Context(), "", []string{runtime.device.ID}) diff --git a/services/core/tests/integration/environment_worker_scan_test.go b/services/core/tests/integration/environment_worker_scan_test.go index 2d255ade0..5149190ff 100644 --- a/services/core/tests/integration/environment_worker_scan_test.go +++ b/services/core/tests/integration/environment_worker_scan_test.go @@ -21,7 +21,7 @@ func TestWorkerEnvironmentRetriesNewlyReadyAtNextScan(t *testing.T) { h.session = publicSession(t, h, "scan-barrier") receipt := h.message("barrier", "ordinary work") scanned := time.Now() - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) // Dispatch starts only after selectWork has examined the pending input on // the same pass, while its exact Runtime is still incapable of preparation. barrier := nextWorkerFrame(t, frames, testExecutionRequest) @@ -69,7 +69,7 @@ func TestWorkerEnvironmentPaginationReachesReadyTail(t *testing.T) { h.session = publicSession(t, h, "page-barrier") receipt := h.message("barrier", "ordinary work") scanned := time.Now() - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) barrier := nextWorkerFrame(t, frames, testExecutionRequest) if barrier.ID != receipt.TurnID { t.Fatal("unexpected page barrier") diff --git a/services/core/tests/integration/environment_worker_test.go b/services/core/tests/integration/environment_worker_test.go index da3aaf6f6..3ba10d5ce 100644 --- a/services/core/tests/integration/environment_worker_test.go +++ b/services/core/tests/integration/environment_worker_test.go @@ -31,7 +31,7 @@ func TestWorkerEnvironmentSharesCapacityThroughClaimAndCleanup(t *testing.T) { receipt := h.message(key, "ordinary") ordinary[receipt.TurnID] = session } - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) var normal []proto.Envelope var preparing []proto.Envelope for range 4 { @@ -132,7 +132,7 @@ func TestWorkerEnvironmentRetriesPendingWithoutExtendingDeadline(t *testing.T) { pending := workerEnvironmentReservation(t, h) runtime := h.environments[pending.SessionID] frames := workerFrames(t, h, runtime) - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) first := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) started := time.Now() handle := acknowledgePreparation(runtime, first.ID) @@ -162,7 +162,7 @@ func TestWorkerEnvironmentRetriesPendingWithoutExtendingDeadline(t *testing.T) { if err != nil || stored.State != sessions.EnvironmentInputPending || !stored.Deadline.Equal(pending.Deadline) || len(stored.Receipts) != 0 { t.Fatal("retry or shutdown changed the original reservation", stored, err) } - _, stop = startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop = startEnvironmentExpiryWorker(t, h.s, h.d) third := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) handle = acknowledgePreparation(runtime, third.ID) runtime.write(third.ID, proto.TypePreparationStatus, proto.PreparationStatusPayload{Handle: handle, Revision: 2, State: "ready"}) diff --git a/services/core/tests/integration/file_resource_semantics_public_test.go b/services/core/tests/integration/file_resource_semantics_public_test.go index c2290b747..83854d136 100644 --- a/services/core/tests/integration/file_resource_semantics_public_test.go +++ b/services/core/tests/integration/file_resource_semantics_public_test.go @@ -32,8 +32,8 @@ func TestFileResourceSemanticsOfficialClientPostgres(t *testing.T) { }) newServer := func() *httptest.Server { t.Helper() - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} - h, err := publicHandler(t, s, db, auth, "codex") + s := NewWithCredentialCipher(pool, cipher) + h, err := publicHandler(t, s, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/fixture_db_test.go b/services/core/tests/integration/fixture_db_test.go deleted file mode 100644 index 5f7361505..000000000 --- a/services/core/tests/integration/fixture_db_test.go +++ /dev/null @@ -1,141 +0,0 @@ -package integration - -import ( - "context" - "errors" - "testing" - - "github.com/jackc/pgx/v5/pgxpool" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/modelconfigurationpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -// fixtureDB is the database and credential key that built the test's Store. -type fixtureDB struct { - pool *pgxpool.Pool - cipher *credentialcrypto.Cipher // nil for a keyless Store - publicURL string // the public URL of the Store's placement rules, if any -} - -// newTestStoreDB is testStore with the fixtureDB that built it. -func newTestStoreDB(t *testing.T) (*Store, fixtureDB) { - s, pool := testStore(t) - return s, fixtureDB{pool: pool} -} - -// newModelTestStoreDB is NewModelTestStore with the fixtureDB that built it. -func newModelTestStoreDB(t *testing.T) (*Store, fixtureDB) { - s, pool := NewModelTestStore(t) - return s, fixtureDB{pool: pool, cipher: fixtureCipher} -} - -// newManagedTestStoreDB is newManagedTestStore with the fixtureDB that built it. -func newManagedTestStoreDB(t *testing.T) (*Store, fixtureDB) { - s, pool := newManagedTestStore(t) - return s, fixtureDB{pool: pool, cipher: fixtureCipher} -} - -// startWorker starts the execution Worker as cmd/server does: it acquires the -// execution lease on db and hands it, with the execution writer built on it, to -// the Worker, which closes it when Run exits. The Worker opens MCP bearer tokens -// through the vaults service on db, records model configuration observations -// through the model configuration adapter on db, runs Session use cases and -// reads through the Session service and adapter on db, and reads the -// deployment through the deployment adapter on db. -func startWorker(t testing.TB, ctx context.Context, db fixtureDB, dispatcher *execution.Dispatcher) *execution.Worker { - t.Helper() - worker, err := startWorkerErr(ctx, db, dispatcher) - if err != nil { - t.Fatal(err) - } - return worker -} - -// startWorkerErr is startWorker for tests that assert a startup failure. -func startWorkerErr(ctx context.Context, db fixtureDB, dispatcher *execution.Dispatcher) (*execution.Worker, error) { - lease, err := pgunit.AcquireLease(ctx, db.pool) - if err != nil { - return nil, err - } - owner, err := fixtureOwner(db, lease) - if err != nil { - return nil, errors.Join(err, lease.Close(ctx)) - } - return startOwnedWorkerErr(ctx, db, dispatcher, owner) -} - -// startOwnedWorker is startWorker on an Owner the test already holds, for tests -// that also run execution operations on it. The Worker closes its lease when -// Run exits. -func startOwnedWorker(t testing.TB, ctx context.Context, db fixtureDB, dispatcher *execution.Dispatcher, owner execution.Owner) *execution.Worker { - t.Helper() - worker, err := startOwnedWorkerErr(ctx, db, dispatcher, owner) - if err != nil { - t.Fatal(err) - } - return worker -} - -func startOwnedWorkerErr(ctx context.Context, db fixtureDB, dispatcher *execution.Dispatcher, owner execution.Owner) (*execution.Worker, error) { - _, credentials, err := fixtureVaults(db) - if err != nil { - return nil, errors.Join(err, owner.Lease.Close(ctx)) - } - sessionStore, sessionService, err := fixtureSessions(db) - if err != nil { - return nil, errors.Join(err, owner.Lease.Close(ctx)) - } - deployments, err := fixtureDeploymentService(db) - if err != nil { - return nil, errors.Join(err, owner.Lease.Close(ctx)) - } - owned := *dispatcher - owned.Credentials = credentials - owned.Observer = modelconfigurationpg.New(pgunit.NewPool(db.pool), db.cipher) - owned.Deployment = deployments - owned.DeploymentReader = deploymentpg.New(pgunit.NewPool(db.pool), db.cipher) - owned.Sessions = sessionService - owned.SessionsReader = sessionStore - return execution.StartWorker(ctx, &owned, owner) -} - -// executionOwner acquires the execution lease on db and builds the execution -// operations on it, for tests that run them without a Worker. The lease -// closes when the test ends. -func executionOwner(t testing.TB, db fixtureDB) execution.Owner { - t.Helper() - lease, err := pgunit.AcquireLease(t.Context(), db.pool) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = lease.Close(context.Background()) }) - owner, err := fixtureOwner(db, lease) - if err != nil { - t.Fatal(err) - } - return owner -} - -// fixtureOwner builds the deployment and Session execution operations on -// lease, as cmd/server does. -func fixtureOwner(db fixtureDB, lease *pgunit.Lease) (execution.Owner, error) { - _, changes, err := fixtureDeploymentExecution(db, lease) - if err != nil { - return execution.Owner{}, err - } - sessionExecution, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) - if err != nil { - return execution.Owner{}, err - } - return execution.Owner{ - Lease: lease, - Deployment: changes, - Sessions: sessionExecution, - }, nil -} diff --git a/services/core/tests/integration/fixtures_test.go b/services/core/tests/integration/fixtures_test.go index c23832821..6fecfab2c 100644 --- a/services/core/tests/integration/fixtures_test.go +++ b/services/core/tests/integration/fixtures_test.go @@ -26,7 +26,7 @@ var fixtureCipher, _ = credentialcrypto.New(bytes.Repeat([]byte{61}, 32)) // can freeze a model provider. func NewModelTestStore(t *testing.T) (*Store, *pgxpool.Pool) { _, pool := testStore(t) - return withPlacement(t, NewWithCredentialCipher(pool, fixtureCipher)), pool + return NewWithCredentialCipher(pool, fixtureCipher), pool } // FixtureModelProvider is a valid bundle for the harness. Hosted and self-hosted diff --git a/services/core/tests/integration/function_execution_native_test.go b/services/core/tests/integration/function_execution_native_test.go index c19732859..96844b3d4 100644 --- a/services/core/tests/integration/function_execution_native_test.go +++ b/services/core/tests/integration/function_execution_native_test.go @@ -18,7 +18,7 @@ func TestNativeFunctionExecutionPersistsCallsResultsAndContinuity(t *testing.T) if err != nil { t.Fatal(err) } - 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) } nativeID := "" @@ -48,11 +48,11 @@ func TestNativeFunctionExecutionPersistsCallsResultsAndContinuity(t *testing.T) } } h.finished(running, sessions.TurnCompleted) - saved, err := FixtureFunctionCall(ctx, h.db.pool, h.tenant, h.session.ID, input.TurnID, action.CallID) + saved, err := FixtureFunctionCall(ctx, h.s.pool, h.tenant, h.session.ID, input.TurnID, action.CallID) if err != nil || !saved.Applied { t.Fatal(saved, err) } - page, err := sessionReads(h.db.pool).ListItems(ctx, h.tenant, h.session.ID, "", 100, true) + page, err := sessionAdapter(h.s).ListItems(ctx, h.tenant, h.session.ID, "", 100, true) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/function_execution_test.go b/services/core/tests/integration/function_execution_test.go index 28dd4ea59..16b67e438 100644 --- a/services/core/tests/integration/function_execution_test.go +++ b/services/core/tests/integration/function_execution_test.go @@ -24,7 +24,7 @@ func newFunctionHarness(t *testing.T) *dispatchHarness { if err != nil { t.Fatal(err) } - 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) } h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, FunctionTools: proto.CapabilitySupported, FunctionResultImages: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) @@ -88,7 +88,7 @@ func TestExecutionFunctionsWaitForEveryApplicationReceipt(t *testing.T) { var reply proto.FunctionResultPayload _ = h.read(proto.TypeFunctionResult).DecodePayload(&reply) public := items.Identity(input.TurnID, "tool:"+reply.CallID) - saved, err := FixtureFunctionCall(t.Context(), h.db.pool, h.tenant, h.session.ID, input.TurnID, public) + saved, err := FixtureFunctionCall(t.Context(), h.s.pool, h.tenant, h.session.ID, input.TurnID, public) if err != nil || saved.Applied || reply.DeliveryID != "function:"+public || len(reply.Content) != 1 || *reply.Content[0].Text != "saved" { t.Fatal(saved, reply, err) } @@ -159,7 +159,7 @@ func TestExecutionFunctionsCancellationAndUnconfirmedResults(t *testing.T) { t.Fatal(bound, err) } } - saved, err := FixtureFunctionCall(t.Context(), h.db.pool, h.tenant, h.session.ID, input.TurnID, id) + saved, err := FixtureFunctionCall(t.Context(), h.s.pool, h.tenant, h.session.ID, input.TurnID, id) if err != nil || saved.Applied || len(saved.Result) == 0 { t.Fatal(saved, err) } @@ -194,7 +194,7 @@ func TestExecutionFunctionsRequireAdvertisedCapability(t *testing.T) { t.Fatal(err) } h.session = session - if err := bindSessionDevice(t, h.db, h.tenant, session.ID, h.device.ID); err != nil { + if err := bindSessionDevice(t, h.s, h.tenant, session.ID, h.device.ID); err != nil { t.Fatal(err) } input := h.message("start", "Run") diff --git a/services/core/tests/integration/function_images_native_test.go b/services/core/tests/integration/function_images_native_test.go index ae4925610..d0c78d6df 100644 --- a/services/core/tests/integration/function_images_native_test.go +++ b/services/core/tests/integration/function_images_native_test.go @@ -31,7 +31,7 @@ func TestNativeFunctionImagePublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) 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() { @@ -47,7 +47,7 @@ func TestNativeFunctionImagePublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(t, worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, auth, kind, workerExecution(t, worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } @@ -72,7 +72,7 @@ func TestNativeFunctionImagePublicExecution(t *testing.T) { t.Fatal("invalid evidence", err) } for _, item := range proof.Calls { - call, err := FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, item.Turn, item.Call) + call, err := FixtureFunctionCall(ctx, h.s.pool, h.tenant, proof.Session, item.Turn, item.Call) if err != nil || !call.Applied { t.Fatal("function delivery acknowledgement missing", err) } diff --git a/services/core/tests/integration/function_input_execution_test.go b/services/core/tests/integration/function_input_execution_test.go index eec056ea2..63dbe8d17 100644 --- a/services/core/tests/integration/function_input_execution_test.go +++ b/services/core/tests/integration/function_input_execution_test.go @@ -52,7 +52,7 @@ func TestExecutionFunctionInputBatchStillSteersMessages(t *testing.T) { } h.write(input.TurnID, proto.TypeDone, proto.DonePayload{Content: "done"}) h.finished(running, sessions.TurnCompleted) - saved, err := FixtureFunctionCall(t.Context(), h.db.pool, h.tenant, h.session.ID, input.TurnID, state.RequiredActions[0].CallID) + saved, err := FixtureFunctionCall(t.Context(), h.s.pool, h.tenant, h.session.ID, input.TurnID, state.RequiredActions[0].CallID) if err != nil || !saved.Applied { t.Fatal(saved, err) } diff --git a/services/core/tests/integration/function_inputs_public_test.go b/services/core/tests/integration/function_inputs_public_test.go index 16b44ec29..486a53276 100644 --- a/services/core/tests/integration/function_inputs_public_test.go +++ b/services/core/tests/integration/function_inputs_public_test.go @@ -21,7 +21,7 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { if python == "" { t.Skip("OAC_TEST_OFFICIAL_SDK_PYTHON is required for official-client verification") } - s, db := newTestStoreDB(t) + s, _ := testStore(t) ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second) defer cancel() tenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() @@ -40,14 +40,14 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { if _, err := transitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } - functions := executionOwner(t, db).Sessions + functions := executionOwner(t, s).Sessions for _, id := range []string{"a", "b", "c", "rollback", "late"} { if err := functions.RecordFunctionCall(ctx, tenant, session.ID, input.TurnID, sessions.FunctionCall{CallID: id, ExecutorCallID: "native-" + id, Name: "lookup", Arguments: json.RawMessage(`{}`)}); err != nil { t.Fatal(err) } } auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}}) - handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) + handler, err := publicHandler(t, s, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } @@ -72,7 +72,7 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { t.Fatalf("SDK admission: %s %v %s", line, err, stderr.String()) } for _, id := range []string{"a", "b", "c", "rollback", "late"} { - call, err := FixtureFunctionCall(ctx, db.pool, tenant, session.ID, input.TurnID, id) + call, err := FixtureFunctionCall(ctx, s.pool, tenant, session.ID, input.TurnID, id) if err != nil || call.Applied { t.Fatal(call, err) } diff --git a/services/core/tests/integration/function_public_native_test.go b/services/core/tests/integration/function_public_native_test.go index 9049edfa7..0da28f9c9 100644 --- a/services/core/tests/integration/function_public_native_test.go +++ b/services/core/tests/integration/function_public_native_test.go @@ -43,7 +43,7 @@ func TestNativePublicFunctionExecution(t *testing.T) { t.Fatal(proof, err) } for i, callID := range proof.Calls { - call, err := FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turns[i], callID) + call, err := FixtureFunctionCall(ctx, h.s.pool, h.tenant, proof.Session, proof.Turns[i], callID) if err != nil || call.Applied != (i < 2) { t.Fatal(call, err) } @@ -60,7 +60,7 @@ func TestNativePublicFunctionExecution(t *testing.T) { func nativePublicFunctionServer(t *testing.T, h *dispatchHarness, ctx context.Context, provider *v1.ModelProviderInput) (string, string) { t.Helper() - worker := startWorker(t, ctx, h.db, h.d) + worker := startWorker(t, ctx, h.s, h.d) ctx, cancel := context.WithCancel(ctx) t.Cleanup(cancel) done := make(chan error, 1) @@ -75,7 +75,7 @@ func nativePublicFunctionServer(t *testing.T, h *dispatchHarness, ctx context.Co }) token := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}}) - handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(t, worker), nativeDeploymentDefaults("gpt-5.5", provider)) + handler, err := publicHandler(t, h.s, auth, "codex", workerExecution(t, worker), nativeDeploymentDefaults("gpt-5.5", provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/function_state_public_test.go b/services/core/tests/integration/function_state_public_test.go index 0b9cedc0a..7c1b84cd7 100644 --- a/services/core/tests/integration/function_state_public_test.go +++ b/services/core/tests/integration/function_state_public_test.go @@ -21,7 +21,7 @@ func TestFunctionStateOfficialClientReadsAndLiveEvents(t *testing.T) { if python == "" { t.Skip("OAC_TEST_OFFICIAL_SDK_PYTHON is required for official-client verification") } - s, db := newTestStoreDB(t) + s, _ := testStore(t) ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second) defer cancel() tenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() @@ -37,7 +37,7 @@ func TestFunctionStateOfficialClientReadsAndLiveEvents(t *testing.T) { if _, err := transitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } - functions := executionOwner(t, db).Sessions + functions := executionOwner(t, s).Sessions record := func(id string) { t.Helper() if err := functions.RecordFunctionCall(ctx, tenant, session.ID, input.TurnID, sessions.FunctionCall{CallID: id, ExecutorCallID: "private-" + id, Name: "lookup", Arguments: json.RawMessage(`{"ticket":9007199254740993}`)}); err != nil { @@ -46,7 +46,7 @@ func TestFunctionStateOfficialClientReadsAndLiveEvents(t *testing.T) { } record("first") auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}}) - handler, err := publicHandler(t, s, db, auth, "codex") + handler, err := publicHandler(t, s, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/function_stream_native_test.go b/services/core/tests/integration/function_stream_native_test.go index 481451e5e..29fcf4954 100644 --- a/services/core/tests/integration/function_stream_native_test.go +++ b/services/core/tests/integration/function_stream_native_test.go @@ -37,7 +37,7 @@ func TestNativePublicFunctionStreamHelper(t *testing.T) { t.Fatal(proof, err) } for i, callID := range proof.Calls { - call, err := FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turns[i], callID) + call, err := FixtureFunctionCall(ctx, h.s.pool, h.tenant, proof.Session, proof.Turns[i], callID) if err != nil || !call.Applied { t.Fatal(call, err) } diff --git a/services/core/tests/integration/function_worker_test.go b/services/core/tests/integration/function_worker_test.go index 5e32cfdde..916f4f113 100644 --- a/services/core/tests/integration/function_worker_test.go +++ b/services/core/tests/integration/function_worker_test.go @@ -42,7 +42,7 @@ func TestWorkerWaitsForToolCapabilities(t *testing.T) { } } if prebound && isMCP { - 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) } } @@ -66,7 +66,7 @@ func TestWorkerWaitsForToolCapabilities(t *testing.T) { input := h.message("queued", "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() { @@ -83,7 +83,7 @@ func TestWorkerWaitsForToolCapabilities(t *testing.T) { t.Fatal(current, 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) } } @@ -125,8 +125,8 @@ func mcpBearerWorkerConfiguration(t *testing.T, h *dispatchHarness) (string, str if err != nil { t.Fatal(err) } - h.s, h.db = NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} - _, service, err := fixtureVaults(h.db) + h.s = NewWithCredentialCipher(pool, cipher) + _, service, err := fixtureVaults(h.s) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/harness_onboarding_test.go b/services/core/tests/integration/harness_onboarding_test.go index dad2531bd..b655f11cb 100644 --- a/services/core/tests/integration/harness_onboarding_test.go +++ b/services/core/tests/integration/harness_onboarding_test.go @@ -49,7 +49,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { started, write, declaration := startOnboardingPeer(t, h) ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker := startWorker(t, ctx, h.db, h.d) + worker := startWorker(t, ctx, h.s, h.d) stopped := make(chan error, 1) go func() { stopped <- worker.Run(ctx) }() defer func() { @@ -62,7 +62,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { }() token := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}}) - handler, err := publicHandler(t, h.s, h.db, auth, "fixture_harness", workerExecution(t, worker), withPolicy(policy)) + handler, err := publicHandler(t, h.s, auth, "fixture_harness", workerExecution(t, worker), withPolicy(policy)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/hosted_initialization_failure_public_test.go b/services/core/tests/integration/hosted_initialization_failure_public_test.go index 15ea8c8ba..ee461f4e5 100644 --- a/services/core/tests/integration/hosted_initialization_failure_public_test.go +++ b/services/core/tests/integration/hosted_initialization_failure_public_test.go @@ -98,19 +98,17 @@ func (p *hostedFailureProvider) prepare(request proto.RuntimePreparePayload, _ [ return completedInitialization(request, nil) } -func hostedFailureStore(t *testing.T) (*Store, fixtureDB) { +func hostedFailureStore(t *testing.T) *Store { t.Helper() _, pool := newManagedTestStore(t) cipher, err := credentialcrypto.New(bytes.Repeat([]byte{7}, 32)) if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} - s.SetPlacement(fixtureRules(t, db)) - return s, db + return NewWithCredentialCipher(pool, cipher) } -func hostedFailureSession(t *testing.T, s *Store, db fixtureDB, tenant string, input sessions.CreateSession) (sessions.Session, sessions.Environment) { +func hostedFailureSession(t *testing.T, s *Store, tenant string, input sessions.CreateSession) (sessions.Session, sessions.Environment) { t.Helper() input.Creator, input.Engine, input.IdempotencyKey = FixtureCreator(), "codex", uuid.NewString() input.Configuration = json.RawMessage(`{"agent":{"id":"agent_test","model":"test-model","tools":[]},"environment":{"type":"openai_hosted","network":{"access":"enabled"}}}`) @@ -121,21 +119,21 @@ func hostedFailureSession(t *testing.T, s *Store, db fixtureDB, tenant string, i if err != nil { t.Fatal(err) } - environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } return session, environment } -func failHostedInitialization(t *testing.T, s *Store, db fixtureDB, tenant string, environment sessions.Environment, p *hostedFailureProvider) { +func failHostedInitialization(t *testing.T, s *Store, tenant string, environment sessions.Environment, p *hostedFailureProvider) { t.Helper() key := uuid.NewString() - w, _ := managedWorkerMode(t, s, db, key, p, false, true) + w, _ := managedWorkerMode(t, s, key, p, false, true) if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err != nil { t.Fatal(err) } - awaitInitialization(t, db.pool, tenant, environment.ID, "failed") + awaitInitialization(t, s, tenant, environment.ID, "failed") } // H1/H2/H3/H4: one transaction records the Environment failure, an error event @@ -182,12 +180,12 @@ func TestHostedInitializationFailureRecordsSafeSessionFailure(t *testing.T) { "Failed to provision environment: Skill installation failed", []string{"configure", "skill"}}, } { t.Run(test.name, func(t *testing.T) { - s, db := hostedFailureStore(t) + s := hostedFailureStore(t) tenant := uuid.NewString() - session, environment := hostedFailureSession(t, s, db, tenant, test.input) + session, environment := hostedFailureSession(t, s, tenant, test.input) p := &hostedFailureProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, fail: test.p.fail, skip: test.p.skip, result: test.p.result, err: test.p.err} - failHostedInitialization(t, s, db, tenant, environment, p) + failHostedInitialization(t, s, tenant, environment, p) if !reflect.DeepEqual(p.steps, test.steps) || p.kills != 0 || p.commandCalls.Load() != 0 { t.Fatal("failed initialization continued or reclaimed compute", p.steps, p.kills) } @@ -247,15 +245,15 @@ func TestHostedInitializationFailureRecordsSafeSessionFailure(t *testing.T) { // A pending initial input settles exactly as before; the one failed snapshot // carries both that settlement and the provisioning failure. func TestHostedInitializationFailureSettlesPendingInitialInput(t *testing.T) { - s, db := hostedFailureStore(t) + s := hostedFailureStore(t) tenant := uuid.NewString() - session, environment := hostedFailureSession(t, s, db, tenant, sessions.CreateSession{ + session, environment := hostedFailureSession(t, s, tenant, sessions.CreateSession{ Initialization: environmentconfig.Setup{Commands: []environmentconfig.SetupCommand{{Command: "exit 3"}}}, InitialInputs: []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"initial"}`)}}, }) p := &hostedFailureProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, fail: "setup", result: failedInitialization(3)} - failHostedInitialization(t, s, db, tenant, environment, p) + failHostedInitialization(t, s, tenant, environment, p) read, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || read.PendingInput || read.EnvironmentInputActivity == nil || read.EnvironmentInputActivity.Status != "failed" || read.EnvironmentInputActivity.Failure != "environment_unavailable" || read.EnvironmentFailure == nil { @@ -281,9 +279,9 @@ func TestHostedInitializationFailureSettlesPendingInitialInput(t *testing.T) { // stream ends after agent.session.failed; later input gets the observed 409; // delete succeeds; tenant B sees nothing; the canary never appears. func TestHostedInitializationFailurePublicHTTP(t *testing.T) { - s, db := hostedFailureStore(t) + s := hostedFailureStore(t) tenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() - session, environment := hostedFailureSession(t, s, db, tenant, sessions.CreateSession{ + session, environment := hostedFailureSession(t, s, tenant, sessions.CreateSession{ Initialization: environmentconfig.Setup{Commands: []environmentconfig.SetupCommand{{Command: "echo " + hostedFailureCanary + "; exit 3"}}}, Metadata: map[string]string{"case": "setup-exit3"}, }) @@ -294,12 +292,12 @@ func TestHostedInitializationFailurePublicHTTP(t *testing.T) { previous := slog.Default() slog.SetDefault(slog.New(slog.NewTextHandler(logs, &slog.HandlerOptions{Level: slog.LevelDebug}))) t.Cleanup(func() { slog.SetDefault(previous) }) - w, _ := managedWorkerMode(t, s, db, key, p, false, true) + w, _ := managedWorkerMode(t, s, key, p, false, true) auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "tenant-b", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, w)) + handler, err := publicHandler(t, s, auth, "codex", workerExecution(t, w)) if err != nil { t.Fatal(err) } @@ -337,7 +335,7 @@ func TestHostedInitializationFailurePublicHTTP(t *testing.T) { if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err != nil { t.Fatal(err) } - awaitInitialization(t, db.pool, tenant, environment.ID, "failed") + awaitInitialization(t, s, tenant, environment.ID, "failed") reason := `Failed to provision environment: script "setup_commands[0]" failed with exit code 3` read, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) diff --git a/services/core/tests/integration/initial_files_http_test.go b/services/core/tests/integration/initial_files_http_test.go index ccf870867..9fee427b4 100644 --- a/services/core/tests/integration/initial_files_http_test.go +++ b/services/core/tests/integration/initial_files_http_test.go @@ -19,12 +19,11 @@ func TestInitialFilesHTTPInlineLimitsAndRetry(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) tenant, token := 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: tenant}}) // Exercise HTTP parsing and durable storage without starting a Runtime. - handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), managedSandboxes(t, db), fixtureDeploymentProvider()) + handler, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), managedSandboxes(t, s), fixtureDeploymentProvider()) if err != nil { t.Fatal(err) } @@ -59,7 +58,7 @@ func TestInitialFilesHTTPInlineLimitsAndRetry(t *testing.T) { id = response.ID } for position := range files { - _, actual, err := fixtureSessionStore(db).ReadInitialEnvironmentFile(t.Context(), tenant, id, position) + _, actual, err := sessionAdapter(s).ReadInitialEnvironmentFile(t.Context(), tenant, id, position) if err != nil || !bytes.Equal(actual, data) { t.Fatal("large HTTP snapshot differs", err) } diff --git a/services/core/tests/integration/input_conflicts_public_test.go b/services/core/tests/integration/input_conflicts_public_test.go index 258ffb39f..41578612f 100644 --- a/services/core/tests/integration/input_conflicts_public_test.go +++ b/services/core/tests/integration/input_conflicts_public_test.go @@ -37,21 +37,21 @@ const ( // and every rejection leaves the database and the pending action unchanged. func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { // An isolated database keeps the no-write digest independent of other tests. - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) ctx := t.Context() tenant, owner, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "conflict-owner", TokenSHA256: runtimedevice.HashCredential(owner), TenantID: tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "conflict-foreign", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) + h, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() client := pathIDClient{t: t, server: server} - functions := executionOwner(t, db).Sessions + functions := executionOwner(t, s).Sessions create := func(environment string, initial bool) string { t.Helper() @@ -133,7 +133,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { // reject checks each response body and that no rejection wrote anything. reject := func(cases []rejection, watched ...string) { t.Helper() - digest := databaseDigest(t, db.pool) + digest := databaseDigest(t, s.pool) before := make([]string, len(watched)) for i, session := range watched { before[i] = read(session) @@ -147,7 +147,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { t.Errorf("%s: %d %s", tc.name, status, body) } } - if after := databaseDigest(t, db.pool); !reflect.DeepEqual(after, digest) { + if after := databaseDigest(t, s.pool); !reflect.DeepEqual(after, digest) { t.Error("rejected input changed the database") } for i, session := range watched { @@ -232,7 +232,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { if status, body := submit(owner, session, "same-after-completion", valid); status != http.StatusAccepted || body != "" { t.Fatalf("identical result after completion: %d %s", status, body) } - if call, err := FixtureFunctionCall(ctx, db.pool, tenant, session, current.TurnID, "pending-call"); err != nil || !bytes.Contains(call.Result, []byte(`"value"`)) { + if call, err := FixtureFunctionCall(ctx, s.pool, tenant, session, current.TurnID, "pending-call"); err != nil || !bytes.Contains(call.Result, []byte(`"value"`)) { t.Fatal("saved result changed", call, err) } @@ -246,7 +246,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { {"after-cancel", owner, session, []string{late}, 409, turnConflictBody, ""}, {"after-cancel-foreign", foreign, session, []string{late}, 404, missingSessionBody, ""}, }, session) - if call, err := FixtureFunctionCall(ctx, db.pool, tenant, session, cancelled.TurnID, "late-call"); err != nil || call.Result != nil { + if call, err := FixtureFunctionCall(ctx, s.pool, tenant, session, cancelled.TurnID, "late-call"); err != nil || call.Result != nil { t.Fatal("late result was saved", call, err) } @@ -272,7 +272,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { if status, body := createSession("first"); status != http.StatusCreated { t.Fatal(status, body) } - digest := databaseDigest(t, db.pool) + digest := databaseDigest(t, s.pool) status, body := createSession("changed") var creation struct { Error struct { @@ -285,7 +285,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { creation.Error.Type != "conflict_error" || creation.Error.Code != "idempotency_conflict" || creation.Error.Param != nil { t.Fatalf("creation key reuse: %d %s", status, body) } - if after := databaseDigest(t, db.pool); !reflect.DeepEqual(after, digest) { + if after := databaseDigest(t, s.pool); !reflect.DeepEqual(after, digest) { t.Error("creation key reuse changed the database") } } diff --git a/services/core/tests/integration/item_order_migration_test.go b/services/core/tests/integration/item_order_migration_test.go index a7ac7c6fc..31b3abf22 100644 --- a/services/core/tests/integration/item_order_migration_test.go +++ b/services/core/tests/integration/item_order_migration_test.go @@ -152,7 +152,7 @@ func TestItemOrderMigrationPreservesIndexedHistory(t *testing.T) { if _, err = transitionTurn(ctx, s, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } - if err = executionOwner(t, fixtureDB{pool: migratedPool}).Sessions.AppendTurnEvents(ctx, tenant, session, turn, 1, []sessions.ExecutionEvent{{Kind: "delta", Payload: json.RawMessage(`{"item_id":"after-upgrade","delta":"continued"}`)}}); err != nil { + if err = executionOwner(t, New(migratedPool)).Sessions.AppendTurnEvents(ctx, tenant, session, turn, 1, []sessions.ExecutionEvent{{Kind: "delta", Payload: json.RawMessage(`{"item_id":"after-upgrade","delta":"continued"}`)}}); err != nil { t.Fatal(err) } addedID := items.Identity(turn, "message:after-upgrade") diff --git a/services/core/tests/integration/item_order_test.go b/services/core/tests/integration/item_order_test.go index 788a2b3bc..d7c152ac7 100644 --- a/services/core/tests/integration/item_order_test.go +++ b/services/core/tests/integration/item_order_test.go @@ -16,7 +16,7 @@ import ( func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) { ctx := context.Background() s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "ordered"}) if err != nil { @@ -30,7 +30,7 @@ func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) if err != nil { t.Fatal(err) } - page, err := sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) + page, err := sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil || len(page.Items) != 1 { t.Fatal(page, err) } @@ -54,7 +54,7 @@ func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) if _, err = sendMessage(ctx, s, tenant, session.ID, "steer", json.RawMessage(`{"text":"continue"}`)); err != nil { t.Fatal(err) } - page, err = sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) + page, err = sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil || len(page.Items) != 5 { t.Fatal(page, err) } @@ -87,7 +87,7 @@ func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) var got []string cursor := "" for { - page, err := sessionReads(pool).ListItems(ctx, tenant, session.ID, cursor, 2, asc) + page, err := sessionAdapter(s).ListItems(ctx, tenant, session.ID, cursor, 2, asc) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/item_reads_test.go b/services/core/tests/integration/item_reads_test.go index 9e15be0c7..31b01fa3b 100644 --- a/services/core/tests/integration/item_reads_test.go +++ b/services/core/tests/integration/item_reads_test.go @@ -15,8 +15,8 @@ import ( func TestItemsRecoverSnapshotsPartialResultsPaginationAndIsolation(t *testing.T) { ctx := context.Background() - s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + s, _ := testStore(t) + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "items"}) if err != nil { @@ -46,7 +46,7 @@ func TestItemsRecoverSnapshotsPartialResultsPaginationAndIsolation(t *testing.T) t.Fatal(err) } } - page, err := sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) + page, err := sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil { t.Fatal(err) } @@ -60,7 +60,7 @@ func TestItemsRecoverSnapshotsPartialResultsPaginationAndIsolation(t *testing.T) if err != nil { t.Fatal(err) } - reopened := sessionReads(pool) + reopened := sessionAdapter(s) page, err = reopened.ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil { t.Fatal(err) @@ -104,13 +104,13 @@ func TestItemsRecoverSnapshotsPartialResultsPaginationAndIsolation(t *testing.T) } other, _ := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "other"}) // A foreign parent is not found before the cursor is read. - if _, err = sessionReads(pool).ListItems(ctx, uuid.NewString(), session.ID, page.Items[0].ID, 20, true); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).ListItems(ctx, uuid.NewString(), session.ID, page.Items[0].ID, 20, true); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } // Another Session's Item is an invalid cursor here, like a missing or malformed one. for _, cursor := range []string{page.Items[0].ID, uuid.NewString(), "not-a-uuid"} { var invalid *sessions.CursorError - if _, err = sessionReads(pool).ListItems(ctx, tenant, other.ID, cursor, 20, true); !errors.As(err, &invalid) || invalid.Message != "Invalid session item ID in `after`" { + if _, err = sessionAdapter(s).ListItems(ctx, tenant, other.ID, cursor, 20, true); !errors.As(err, &invalid) || invalid.Message != "Invalid session item ID in `after`" { t.Fatal(cursor, err) } } @@ -118,8 +118,8 @@ func TestItemsRecoverSnapshotsPartialResultsPaginationAndIsolation(t *testing.T) func TestItemProjectionFailureRollsBackJournalAndAggregateRecovers(t *testing.T) { ctx := context.Background() - s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + s, _ := testStore(t) + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, _ := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "legacy"}) input, err := sendMessage(ctx, s, tenant, session.ID, "input", json.RawMessage(`{"text":"test"}`)) @@ -138,7 +138,7 @@ func TestItemProjectionFailureRollsBackJournalAndAggregateRecovers(t *testing.T) if err != nil || len(events) != 0 { t.Fatal(events, err) } - page, err := sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) + page, err := sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil || len(page.Items) != 1 { t.Fatal(page, err) } @@ -146,7 +146,7 @@ func TestItemProjectionFailureRollsBackJournalAndAggregateRecovers(t *testing.T) if err != nil { t.Fatal(err) } - current, err := sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) + current, err := sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil || len(current.Items) != 2 || *current.Items[1].Content[0].Text != "legacy answer" { t.Fatal(current, err) } @@ -154,8 +154,8 @@ func TestItemProjectionFailureRollsBackJournalAndAggregateRecovers(t *testing.T) func TestReceiptOnlyTextRecoversWithoutInventingCompletion(t *testing.T) { ctx := context.Background() - s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + s, _ := testStore(t) + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() for _, receiptOnly := range []bool{true, false} { session, _ := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString()}) @@ -177,7 +177,7 @@ func TestReceiptOnlyTextRecoversWithoutInventingCompletion(t *testing.T) { if err != nil { t.Fatal(err) } - page, err := sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) + page, err := sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil || len(page.Items) != 2 || page.Items[1].Status != "incomplete" || *page.Items[1].Content[0].Text != "retained cancellation text" { t.Fatal(page, err) } @@ -186,8 +186,8 @@ func TestReceiptOnlyTextRecoversWithoutInventingCompletion(t *testing.T) { func TestLegacyFailureRetainsPartialAnswerAcrossRecovery(t *testing.T) { ctx := context.Background() - s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + s, _ := testStore(t) + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "failed-items"}) if err != nil { @@ -215,7 +215,7 @@ func TestLegacyFailureRetainsPartialAnswerAcrossRecovery(t *testing.T) { if err != nil { t.Fatal(err) } - page, err := sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) + page, err := sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil || len(page.Items) != 4 { t.Fatal(page, err) } diff --git a/services/core/tests/integration/list_cursor_public_test.go b/services/core/tests/integration/list_cursor_public_test.go index 29d801bd7..9cf08ae1a 100644 --- a/services/core/tests/integration/list_cursor_public_test.go +++ b/services/core/tests/integration/list_cursor_public_test.go @@ -226,24 +226,23 @@ func TestListCursorErrorsPostgres(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) owner, foreign := uuid.NewString(), uuid.NewString() ownerTenant, foreignTenant := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "cursor-owner", TokenSHA256: runtimedevice.HashCredential(owner), TenantID: ownerTenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "cursor-foreign", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: foreignTenant}, }) - 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() client := pathIDClient{t: t, server: server} - leased := executionOwner(t, db) - skillService := SkillService(t, db.pool, db.cipher) - _, sessionService, err := fixtureSessions(db) + leased := executionOwner(t, s) + skillService := SkillService(t, s.pool, s.credentialCipher) + sessionService, err := newSessionService(s) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/list_query_public_test.go b/services/core/tests/integration/list_query_public_test.go index 059d7673e..8a67221f8 100644 --- a/services/core/tests/integration/list_query_public_test.go +++ b/services/core/tests/integration/list_query_public_test.go @@ -25,7 +25,7 @@ func TestListQueryOfficialClientPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) token, foreign := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "query-owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: uuid.NewString()}, @@ -33,7 +33,7 @@ func TestListQueryOfficialClientPostgres(t *testing.T) { }) // Use real admission while leaving dispatch paused. Public cancellation retains // the queued history; this fixture does not perform native or model execution. - worker := startWorker(t, t.Context(), db, &execution.Dispatcher{}) + worker := startWorker(t, t.Context(), s, &execution.Dispatcher{}) t.Cleanup(func() { stopped, cancel := context.WithCancel(context.Background()) cancel() @@ -41,7 +41,7 @@ func TestListQueryOfficialClientPostgres(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, worker)) + handler, err := publicHandler(t, s, auth, "codex", workerExecution(t, worker)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/local_artifact_export_test.go b/services/core/tests/integration/local_artifact_export_test.go index abd071810..06775a674 100644 --- a/services/core/tests/integration/local_artifact_export_test.go +++ b/services/core/tests/integration/local_artifact_export_test.go @@ -41,7 +41,7 @@ func completeLocalArtifactExport(t *testing.T, h *dispatchHarness, worker *execu if next.DecodePayload(&export) != nil || export.Step != "next" || export.Offset != int64(data.Len()) { t.Fatal("export did not await native completion") } - page, err := sessionReads(h.db.pool).ListSessionArtifacts(t.Context(), h.tenant, h.session.ID, "", "", 20, false) + page, err := sessionAdapter(h.s).ListSessionArtifacts(t.Context(), h.tenant, h.session.ID, "", "", 20, false) if err != nil || len(page.Artifacts) != 0 { t.Fatal("capture published before native completion", err) } diff --git a/services/core/tests/integration/local_environment_file_write_test.go b/services/core/tests/integration/local_environment_file_write_test.go index b256291d0..0b802a30f 100644 --- a/services/core/tests/integration/local_environment_file_write_test.go +++ b/services/core/tests/integration/local_environment_file_write_test.go @@ -52,7 +52,7 @@ func TestLocalEnvironmentFileWriteOwnsMutationBeforeDispatch(t *testing.T) { if begin.DecodePayload(&request) != nil || request.Step != "begin" || request.EnvironmentID != environment.ID || request.SessionID != h.session.ID || request.Path != "input" || request.SizeBytes != 3 { t.Fatal("upload identity changed") } - intent, err := FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, begin.ID) + intent, err := FixtureFileWrite(t.Context(), h.s.pool, h.tenant, environment.ID, begin.ID) if err != nil || intent.State != "pending" || intent.Identity.DeviceID != h.device.ID { t.Fatal("dispatch preceded durable ownership", intent, err) } @@ -75,7 +75,7 @@ func TestLocalEnvironmentFileWriteOwnsMutationBeforeDispatch(t *testing.T) { } h.write(begin.ID, proto.TypeWorkspaceWriteResult, proto.WorkspaceWriteResultPayload{Outcome: "completed", SizeBytes: 3}) awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "detached durable commit", func() bool { - got, e := FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, begin.ID) + got, e := FixtureFileWrite(t.Context(), h.s.pool, h.tenant, environment.ID, begin.ID) return e == nil && got.State == "committed" }) session, err := sessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) @@ -94,7 +94,7 @@ func TestLocalEnvironmentFileWriteLostReceiptRemainsPending(t *testing.T) { if result := awaitLocalWrite(t, done); !errors.Is(result.err, execution.ErrExecutionUnavailable) { t.Fatal(result.err) } - intent, err := FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, begin.ID) + intent, err := FixtureFileWrite(t.Context(), h.s.pool, h.tenant, environment.ID, begin.ID) if err != nil || intent.State != "pending" { t.Fatal("disconnect guessed rejection", intent, err) } @@ -112,7 +112,7 @@ func TestLocalEnvironmentFileWriteKnownRejectionReleasesMutation(t *testing.T) { if result := awaitLocalWrite(t, done); !errors.Is(result.err, execution.ErrExecutionUnavailable) { t.Fatal(result.err) } - intent, err := FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, begin.ID) + intent, err := FixtureFileWrite(t.Context(), h.s.pool, h.tenant, environment.ID, begin.ID) if err != nil || intent.State != "rejected" { t.Fatal("rejection did not settle", intent, err) } diff --git a/services/core/tests/integration/local_environment_worker_test.go b/services/core/tests/integration/local_environment_worker_test.go index 76740497a..494efc624 100644 --- a/services/core/tests/integration/local_environment_worker_test.go +++ b/services/core/tests/integration/local_environment_worker_test.go @@ -21,7 +21,7 @@ func localWorker(t *testing.T, scoped, execute bool) (*dispatchHarness, *executi _, pool := testStore(t) insertWorkerRuntimeAllocation(t, pool, h, "disabled") } - environment, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) + environment, err := sessionAdapter(h.s).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) if err != nil { t.Fatal(err) } @@ -41,7 +41,7 @@ func localWorker(t *testing.T, scoped, execute bool) (*dispatchHarness, *executi info, _, _ := peer.AgentKindStatus("codex") return info.Capabilities.LocalEnvironment }) - w := startWorker(t, t.Context(), h.db, h.d) + w := startWorker(t, t.Context(), h.s, h.d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- w.Run(ctx) }() @@ -95,14 +95,14 @@ func TestLocalEnvironmentWorkerRejectsGeneralDeviceDespiteCapability(t *testing. if err != nil { t.Fatal(err) } - unassigned, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, other.ID) + unassigned, err := sessionAdapter(h.s).GetSessionEnvironment(t.Context(), h.tenant, other.ID) if err != nil { t.Fatal(err) } if _, err := w.ReadEnvironmentDirectory(t.Context(), unassigned, "reports"); !errors.Is(err, execution.ErrExecutionUnavailable) { t.Fatal("unassigned environment selected general device", err) } - if _, err := fixtureSessionStore(h.db).GetSessionDevice(t.Context(), h.tenant, other.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(h.s).GetSessionDevice(t.Context(), h.tenant, other.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("read persisted an unauthorized placement", err) } } @@ -151,7 +151,7 @@ func TestLocalEnvironmentWorkerSchedulesPreparationWithoutRemoteResolver(t *test if err != nil || bound.Device.EnvironmentID != environment.ID || bound.NativeSessionID != "local-native-history" { t.Fatal("local native identity was not retained", err) } - artifacts, err := sessionReads(h.db.pool).ListSessionArtifacts(t.Context(), h.tenant, h.session.ID, environment.ID, "", 20, false) + artifacts, err := sessionAdapter(h.s).ListSessionArtifacts(t.Context(), h.tenant, h.session.ID, environment.ID, "", 20, false) if err != nil || len(artifacts.Artifacts) != 1 || artifacts.Artifacts[0].Path != "/workspace/outputs/result.bin" || artifacts.Artifacts[0].TurnID != start.RunID || artifacts.Artifacts[0].SizeBytes != 3 { t.Fatalf("completed turn did not publish output: %+v %v", artifacts, err) } diff --git a/services/core/tests/integration/managed_fixture_test.go b/services/core/tests/integration/managed_fixture_test.go index 2fb5b7ff6..27cac8f77 100644 --- a/services/core/tests/integration/managed_fixture_test.go +++ b/services/core/tests/integration/managed_fixture_test.go @@ -14,5 +14,5 @@ func newManagedTestStore(t *testing.T) (*Store, *pgxpool.Pool) { t.Helper() pool := pgtest.OpenIsolated(t, nil) // Hosted Sessions freeze a model provider, which needs a credential key. - return withPlacement(t, NewWithCredentialCipher(pool, fixtureCipher)), pool + return NewWithCredentialCipher(pool, fixtureCipher), pool } diff --git a/services/core/tests/integration/mcode_public_native_test.go b/services/core/tests/integration/mcode_public_native_test.go index 03198bdcc..85ed9fca6 100644 --- a/services/core/tests/integration/mcode_public_native_test.go +++ b/services/core/tests/integration/mcode_public_native_test.go @@ -31,7 +31,7 @@ func TestNativeMCodePublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 12*time.Minute) defer cancel() - worker := startWorker(t, ctx, h.db, h.d) + worker := startWorker(t, ctx, h.s, h.d) stopped := make(chan error, 1) go func() { stopped <- worker.Run(ctx) }() defer func() { @@ -47,7 +47,7 @@ func TestNativeMCodePublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, "mcode", workerExecution(t, worker), acceptUnavailable(t), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, auth, "mcode", workerExecution(t, worker), acceptUnavailable(t), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/mcp_credential_selection_public_test.go b/services/core/tests/integration/mcp_credential_selection_public_test.go index 72bab89e5..b047c7608 100644 --- a/services/core/tests/integration/mcp_credential_selection_public_test.go +++ b/services/core/tests/integration/mcp_credential_selection_public_test.go @@ -30,13 +30,13 @@ func TestMCPCredentialSelectionPublicPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) tenantA, tokenA, tokenB := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "selection-a", TokenSHA256: runtimedevice.HashCredential(tokenA), TenantID: tenantA}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "selection-b", TokenSHA256: runtimedevice.HashCredential(tokenB), 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) } diff --git a/services/core/tests/integration/message_image_admission_test.go b/services/core/tests/integration/message_image_admission_test.go index d68af03d6..02024727a 100644 --- a/services/core/tests/integration/message_image_admission_test.go +++ b/services/core/tests/integration/message_image_admission_test.go @@ -28,7 +28,7 @@ func TestUnqualifiedImageAdmissionIsAtomic(t *testing.T) { profile, _ := (engine.Catalog{}).Lookup("codex") profile.MessageImages = proto.CapabilityUnsupported h.d.Policy = execution.Policy{Engines: engine.NewCatalog(map[string]engine.Profile{"codex": profile})} - worker := startWorker(t, t.Context(), h.db, h.d) + worker := startWorker(t, t.Context(), h.s, h.d) defer func() { ctx, cancel := context.WithCancel(context.Background()); cancel(); _ = worker.Run(ctx) }() configuration := json.RawMessage(`{"agent":{"model":"fixture"},"environment":{"type":"` + placement + `","workspace_directory":"/workspace"}}`) create := sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "image-create", Configuration: configuration, InitialInputs: imageAdmissionBatch()} diff --git a/services/core/tests/integration/message_images_native_test.go b/services/core/tests/integration/message_images_native_test.go index 5e1a20d4d..d6500f390 100644 --- a/services/core/tests/integration/message_images_native_test.go +++ b/services/core/tests/integration/message_images_native_test.go @@ -32,7 +32,7 @@ func TestNativeMessageImagePublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) 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() { @@ -48,7 +48,7 @@ func TestNativeMessageImagePublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(t, worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, auth, kind, workerExecution(t, worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } @@ -73,7 +73,7 @@ func TestNativeMessageImagePublicExecution(t *testing.T) { if err != nil || json.Unmarshal(raw, &proof) != nil { t.Fatal("invalid evidence", err) } - call, err := FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turn, proof.Call) + call, err := FixtureFunctionCall(ctx, h.s.pool, h.tenant, proof.Session, proof.Turn, proof.Call) if err != nil || !call.Applied { t.Fatal("function application receipt missing", err) } diff --git a/services/core/tests/integration/model_protocol_native_test.go b/services/core/tests/integration/model_protocol_native_test.go index 66b231bdd..587074d0c 100644 --- a/services/core/tests/integration/model_protocol_native_test.go +++ b/services/core/tests/integration/model_protocol_native_test.go @@ -61,7 +61,7 @@ func TestNativeModelProtocolPublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 15*time.Minute) defer cancel() - worker, err := startWorkerErr(ctx, h.db, h.d) + worker, err := startWorkerErr(ctx, h.s, h.d) if err != nil { t.Fatal("cannot start native execution worker") } @@ -78,7 +78,7 @@ func TestNativeModelProtocolPublicExecution(t *testing.T) { token := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}}) providerRevision := uuid.New() - handler, err := publicHandler(t, h.s, h.db, auth, options.Engine, workerExecution(t, worker), withPolicy(h.d.Policy), modelProviderDefaults(func(context.Context, string) (*modelconfiguration.Snapshot, error) { + handler, err := publicHandler(t, h.s, auth, options.Engine, workerExecution(t, worker), withPolicy(h.d.Policy), modelProviderDefaults(func(context.Context, string) (*modelconfiguration.Snapshot, error) { return &modelconfiguration.Snapshot{Model: options.Model, HarnessConfig: options.HarnessConfig, Provider: &options.Provider, Revision: providerRevision}, nil })) if err != nil { @@ -120,7 +120,7 @@ func TestNativeModelProtocolPublicExecution(t *testing.T) { } failed := 0 for _, item := range proof.Calls { - call, err := FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, item.Turn, item.Call) + call, err := FixtureFunctionCall(ctx, h.s.pool, h.tenant, proof.Session, item.Turn, item.Call) if err != nil || !call.Applied { t.Fatal("public function result lacks native delivery acknowledgement") } diff --git a/services/core/tests/integration/native_environment_test.go b/services/core/tests/integration/native_environment_test.go index 8614a7637..dc02c61a2 100644 --- a/services/core/tests/integration/native_environment_test.go +++ b/services/core/tests/integration/native_environment_test.go @@ -102,7 +102,7 @@ func TestNativeNoExecutionEnvironment(t *testing.T) { if err != nil { t.Fatal(err) } - 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) } first := h.message("first", "Return an answer.") @@ -123,7 +123,7 @@ func TestNativeNoExecutionEnvironment(t *testing.T) { if _, err := os.Stat(marker); !os.IsNotExist(err) { t.Fatalf("forbidden command may have executed: %v", err) } - page, err := sessionReads(h.db.pool).ListItems(ctx, h.tenant, h.session.ID, "", 100, true) + page, err := sessionAdapter(h.s).ListItems(ctx, h.tenant, h.session.ID, "", 100, true) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/native_public_execution_test.go b/services/core/tests/integration/native_public_execution_test.go index b370d5ad6..9fbf1e266 100644 --- a/services/core/tests/integration/native_public_execution_test.go +++ b/services/core/tests/integration/native_public_execution_test.go @@ -25,7 +25,7 @@ func verifyNativePublicExecution(t *testing.T, h *dispatchHarness, parent contex defer cancel() // The Turns before this proof ran on the harness's execution Owner, so the // Worker takes that lease rather than a second one. - 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() { @@ -38,7 +38,7 @@ func verifyNativePublicExecution(t *testing.T, h *dispatchHarness, parent contex }() 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: h.tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}}) - handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(t, worker), nativeDeploymentDefaults("gpt-5.5", provider)) + handler, err := publicHandler(t, h.s, auth, "codex", workerExecution(t, worker), nativeDeploymentDefaults("gpt-5.5", provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/no_environment_test.go b/services/core/tests/integration/no_environment_test.go index 2c19405b5..fe6512490 100644 --- a/services/core/tests/integration/no_environment_test.go +++ b/services/core/tests/integration/no_environment_test.go @@ -15,7 +15,7 @@ func TestNoEnvironmentRejectsUnadvertisedDeviceBeforeClaim(t *testing.T) { if err != nil { t.Fatal(err) } - 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) } input := h.message("first", "Answer") diff --git a/services/core/tests/integration/path_id_semantics_public_test.go b/services/core/tests/integration/path_id_semantics_public_test.go index afab3adf7..ebb0b24df 100644 --- a/services/core/tests/integration/path_id_semantics_public_test.go +++ b/services/core/tests/integration/path_id_semantics_public_test.go @@ -91,14 +91,14 @@ func TestMalformedPathIDsMatchMissingPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) owner, foreign := uuid.NewString(), uuid.NewString() ownerTenant := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "path-owner", TokenSHA256: runtimedevice.HashCredential(owner), TenantID: ownerTenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "path-foreign", 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) } @@ -360,7 +360,7 @@ func TestMalformedPathIDsMatchMissingPostgres(t *testing.T) { } // Storage availability checks also run before the lookup of a missing identifier. - h, err = publicHandler(t, New(pool), fixtureDB{pool: pool}, auth, "codex") + h, err = publicHandler(t, New(pool), auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/prepared_dispatch_test.go b/services/core/tests/integration/prepared_dispatch_test.go index 9b2c5ff61..8a01bb556 100644 --- a/services/core/tests/integration/prepared_dispatch_test.go +++ b/services/core/tests/integration/prepared_dispatch_test.go @@ -84,7 +84,7 @@ func TestPreparedDispatchPromotesOriginalBatchAndPersistsCompletion(t *testing.T if err != nil || session.LastTurn != nil { t.Fatal("preparation created work before readiness", session, err) } - items, err := sessionReads(h.db.pool).ListItems(t.Context(), h.tenant, h.session.ID, "", 100, true) + items, err := sessionAdapter(h.s).ListItems(t.Context(), h.tenant, h.session.ID, "", 100, true) if err != nil || len(items.Items) != 0 { t.Fatal("preparation published input history", items, err) } diff --git a/services/core/tests/integration/project_api_keys_http_test.go b/services/core/tests/integration/project_api_keys_http_test.go index 36c3d7014..0a6871812 100644 --- a/services/core/tests/integration/project_api_keys_http_test.go +++ b/services/core/tests/integration/project_api_keys_http_test.go @@ -13,13 +13,13 @@ import ( ) func TestProjectAndSharedKeysHTTPManagement(t *testing.T) { - st, db := newTestStoreDB(t) + st, _ := testStore(t) adminToken := uuid.NewString() admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential(adminToken)}) if err != nil { t.Fatal(err) } - h, err := publicHandler(t, st, db, nil, "codex", storeKeys(st), withCoreKeys(admin)) + h, err := publicHandler(t, st, nil, "codex", storeKeys(st), withCoreKeys(admin)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/projects_public_fixture_test.go b/services/core/tests/integration/projects_public_fixture_test.go index 9fac07967..50f1265e0 100644 --- a/services/core/tests/integration/projects_public_fixture_test.go +++ b/services/core/tests/integration/projects_public_fixture_test.go @@ -10,10 +10,10 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" ) -// fixtureProjects builds the Project adapter and service on db. -func fixtureProjects(t testing.TB, db fixtureDB) (*projectpg.Store, *projects.Service) { +// fixtureProjects builds the Project adapter and service on s. +func fixtureProjects(t testing.TB, s *Store) (*projectpg.Store, *projects.Service) { t.Helper() - projectStore := projectpg.New(pgunit.NewPool(db.pool)) + projectStore := projectpg.New(pgunit.NewPool(s.pool)) projectService, err := projects.NewService(projectStore) if err != nil { t.Fatal(err) diff --git a/services/core/tests/integration/public_execution_test.go b/services/core/tests/integration/public_execution_test.go index 4a45ff9aa..879c09a1d 100644 --- a/services/core/tests/integration/public_execution_test.go +++ b/services/core/tests/integration/public_execution_test.go @@ -27,7 +27,7 @@ func TestExecutionWorkerAdmissionBindingAndRecovery(t *testing.T) { h.session = publicSession(t, h, "public") ctx, cancel := context.WithCancel(context.Background()) 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) }() t.Cleanup(func() { @@ -38,7 +38,7 @@ func TestExecutionWorkerAdmissionBindingAndRecovery(t *testing.T) { t.Error("worker did not stop") } }) - if second, err := startWorkerErr(ctx, h.db, h.d); err == nil { + if second, err := startWorkerErr(ctx, h.s, h.d); err == nil { cancel() go second.Run(ctx) t.Fatal("second service acquired database") @@ -63,7 +63,7 @@ func TestExecutionWorkerAdmissionBindingAndRecovery(t *testing.T) { if inputTextForTest(t, prompt.Input) != "First\n\nSecond" || !prompt.DisableExecutionEnvironment || !prompt.DisableSubagents || prompt.ExecutionControls == nil || *prompt.ExecutionControls != (proto.ExecutionControls{WebSearch: "disabled", TextVerbosity: "medium"}) || prompt.AgentOptions["web_search"] != nil || prompt.AgentOptions["model_verbosity"] != nil { t.Fatal(prompt) } - bound, err := fixtureSessionStore(h.db).GetSessionDevice(ctx, h.tenant, h.session.ID) + bound, err := sessionAdapter(h.s).GetSessionDevice(ctx, h.tenant, h.session.ID) if err != nil || bound.ID != h.device.ID { t.Fatal(bound, err) } @@ -73,7 +73,7 @@ func TestExecutionWorkerAdmissionBindingAndRecovery(t *testing.T) { } h.write(request.ID, proto.TypeDone, proto.DonePayload{Content: "Answer", Metadata: map[string]any{proto.DoneMetaAgentSessionID: "worker-native"}}) waitTurn(t, h, receipts[0].TurnID, sessions.TurnCompleted) - items, err := sessionReads(h.db.pool).ListItems(ctx, h.tenant, h.session.ID, "", 100, true) + items, err := sessionAdapter(h.s).ListItems(ctx, h.tenant, h.session.ID, "", 100, true) if err != nil || len(items.Items) != 3 { t.Fatal(items, err) } @@ -147,7 +147,7 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { if _, err := sendMessage(ctx, h.s, h.tenant, queued.ID, "first", json.RawMessage(`{"text":"Not sent"}`)); err != nil { t.Fatal(err) } - worker := startOwnedWorker(t, ctx, h.db, h.d, h.owner()) + worker := startOwnedWorker(t, ctx, h.s, h.d, h.owner()) stopped, cancel := context.WithCancel(ctx) cancel() if err := worker.Run(stopped); err != context.Canceled { @@ -168,7 +168,7 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { if err != nil || pending.LastTurn.Status != sessions.TurnCancelled { t.Fatal(pending, err) } - restarted, err := startWorkerErr(ctx, h.db, h.d) + restarted, err := startWorkerErr(ctx, h.s, h.d) if err != nil { t.Fatal("lease not released", err) } diff --git a/services/core/tests/integration/public_handler_fixture_test.go b/services/core/tests/integration/public_handler_fixture_test.go index 5c299fcf8..382dfec57 100644 --- a/services/core/tests/integration/public_handler_fixture_test.go +++ b/services/core/tests/integration/public_handler_fixture_test.go @@ -34,49 +34,48 @@ import ( // sets another with executorURL. const testExecutorURL = "wss://core.example/api/v1/agent-daemon/ws" -// publicHandler serves s through api.NewHandler. s backs every area the Store -// implements, and db is the database and credential key that built s; the -// audit reads, Agents, Files, Vaults, Environment reads, Items, Subagents and -// Artifacts come from db. keys authenticate as Project keys and "admin" as the -// Core key. Metrics, Runtime observation and history, and executor connections -// are strict stand-ins. Execution and Sandboxes stay disabled unless configure -// sets them. -func publicHandler(t testing.TB, s *Store, db fixtureDB, keys fixtureKeyResolver, engine string, configure ...func(*api.Dependencies)) (http.Handler, error) { +// publicHandler serves s through api.NewHandler, with every area built on s's +// database, credential key and placement rules as cmd/server builds it. keys +// authenticate as Project keys and "admin" as the Core key. Metrics, Runtime +// observation and history, and executor connections are strict stand-ins. +// Execution and Sandboxes stay disabled unless configure sets them. +func publicHandler(t testing.TB, s *Store, keys fixtureKeyResolver, engine string, configure ...func(*api.Dependencies)) (http.Handler, error) { t.Helper() admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential("admin")}) if err != nil { return nil, err } strict := strictStandIn{t} - audit := auditpg.New(pgunit.NewPool(db.pool)) - agentStore, agentService := fixtureAgents(t, db) - fileStore, fileService := fixtureFiles(t, db) - vaultStore, vaultService, err := fixtureVaults(db) + audit := auditpg.New(pgunit.NewPool(s.pool)) + agentStore, agentService := fixtureAgents(t, s) + fileStore, fileService := fixtureFiles(t, s) + vaultStore, vaultService, err := fixtureVaults(s) if err != nil { return nil, err } - templates := templatepg.New(pgunit.NewPool(db.pool), db.cipher) + templates := templatepg.New(pgunit.NewPool(s.pool), s.credentialCipher) environmentTemplates, err := environmenttemplates.NewService(templates) if err != nil { return nil, err } - modelConfigurationStore := modelconfigurationpg.New(pgunit.NewPool(db.pool), db.cipher) + modelConfigurationStore := modelconfigurationpg.New(pgunit.NewPool(s.pool), s.credentialCipher) modelConfigurationService, err := modelconfiguration.NewService(modelConfigurationStore) if err != nil { return nil, err } - skillStore := skillpg.New(pgunit.NewPool(db.pool), db.cipher) + skillStore := skillpg.New(pgunit.NewPool(s.pool), s.credentialCipher) skillService, err := skills.NewService(skillStore, skillStore) if err != nil { return nil, err } - projectStore, projectService := fixtureProjects(t, db) - sessionStore, service, err := fixtureSessions(db) + projectStore, projectService := fixtureProjects(t, s) + sessionStore := sessionAdapter(s) + service, err := newSessionService(s) if err != nil { return nil, err } deps := api.Dependencies{ - Engine: engine, CoreKeys: admin, InstallationBindings: fixtureDeployment(t, db), + Engine: engine, CoreKeys: admin, InstallationBindings: deploymentService(t, s), Projects: projectService, ProjectsReader: fixtureProjectsReader{Reader: projectStore, keys: keys}, ModelProviders: modelConfigurationService, ModelProvidersReader: modelConfigurationStore, Vaults: vaultService, VaultsReader: vaultStore, @@ -86,7 +85,7 @@ func publicHandler(t testing.TB, s *Store, db fixtureDB, keys fixtureKeyResolver Agents: agentService, AgentsReader: agentStore, Sessions: service, SessionsReader: sessionStore, - SessionCreation: sessionService(t, s), + SessionCreation: service, SessionEvents: sessionStore, Turns: sessionStore, Items: sessionStore, @@ -102,10 +101,10 @@ func publicHandler(t testing.TB, s *Store, db fixtureDB, keys fixtureKeyResolver return api.NewHandler(deps) } -// fixtureAgents builds the Agent adapter and service on db. -func fixtureAgents(t testing.TB, db fixtureDB) (*agentpg.Store, *agents.Service) { +// fixtureAgents builds the Agent adapter and service on s. +func fixtureAgents(t testing.TB, s *Store) (*agentpg.Store, *agents.Service) { t.Helper() - agentStore := agentpg.New(pgunit.NewPool(db.pool), db.cipher) + agentStore := agentpg.New(pgunit.NewPool(s.pool), s.credentialCipher) agentService, err := agents.NewService(agentStore) if err != nil { t.Fatal(err) @@ -113,10 +112,10 @@ func fixtureAgents(t testing.TB, db fixtureDB) (*agentpg.Store, *agents.Service) return agentStore, agentService } -// fixtureFiles builds the File adapter and service on db. -func fixtureFiles(t testing.TB, db fixtureDB) (*filepg.Store, *files.Service) { +// fixtureFiles builds the File adapter and service on s. +func fixtureFiles(t testing.TB, s *Store) (*filepg.Store, *files.Service) { t.Helper() - fileStore := filepg.New(pgunit.NewPool(db.pool)) + fileStore := filepg.New(pgunit.NewPool(s.pool)) fileService, err := files.NewService(fileStore) if err != nil { t.Fatal(err) @@ -179,15 +178,15 @@ func executorURL(url string) func(*api.Dependencies) { return func(d *api.Dependencies) { d.Execution.ExecutorURL = url } } -// managedSandboxes enables the managed sandbox deployment on db: its +// managedSandboxes enables the managed sandbox deployment on s: its // administration and node routes and openai_hosted Environments. Deployment // changes, reset and discovery need the Worker and are strict stand-ins. It // follows the option that enables Execution. -func managedSandboxes(t testing.TB, db fixtureDB) func(*api.Dependencies) { +func managedSandboxes(t testing.TB, s *Store) func(*api.Dependencies) { return func(d *api.Dependencies) { d.Sandboxes = &api.Sandboxes{ - Deployment: fixtureDeployment(t, db), - NodeAllocations: fixtureReader(db), + Deployment: deploymentService(t, s), + NodeAllocations: deploymentStore(s), DeploymentChanges: strictStandIn{t}, DeploymentReset: strictStandIn{t}, ConfigurationDiscovery: strictStandIn{t}, diff --git a/services/core/tests/integration/remote_mcp_credentials_test.go b/services/core/tests/integration/remote_mcp_credentials_test.go index 0a5624583..393e50c9b 100644 --- a/services/core/tests/integration/remote_mcp_credentials_test.go +++ b/services/core/tests/integration/remote_mcp_credentials_test.go @@ -13,12 +13,12 @@ import ( func TestSelfHostedServiceMCPRejectionDoesNotRequireCredentialDecryption(t *testing.T) { for _, mode := range []string{"missing key", "deleted", "tampered"} { t.Run(mode, func(t *testing.T) { - s, db, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) + s, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) switch mode { case "missing key": - s, db = New(db.pool), fixtureDB{pool: db.pool} + s = New(s.pool) case "deleted": - _, service, err := fixtureVaults(db) + _, service, err := fixtureVaults(s) if err != nil { t.Fatal(err) } @@ -26,11 +26,11 @@ func TestSelfHostedServiceMCPRejectionDoesNotRequireCredentialDecryption(t *test t.Fatal(err) } case "tampered": - if _, err := db.pool.Exec(t.Context(), "UPDATE vault_credentials SET token_ciphertext=set_byte(token_ciphertext,15,get_byte(token_ciphertext,15) # 1) WHERE id=$1", credential.ID); err != nil { + if _, err := s.pool.Exec(t.Context(), "UPDATE vault_credentials SET token_ciphertext=set_byte(token_ciphertext,15,get_byte(token_ciphertext,15) # 1) WHERE id=$1", credential.ID); err != nil { t.Fatal(err) } } - handler := selfHostedMCPAdmissionHandler(t, s, db, tenant) + handler := selfHostedMCPAdmissionHandler(t, s, tenant) for _, initial := range []bool{false, true} { body := map[string]any{ "agent": map[string]any{"model": "model", "tools": []any{map[string]any{ @@ -64,7 +64,7 @@ func TestSelfHostedServiceMCPRejectionDoesNotRequireCredentialDecryption(t *test if !strings.Contains(response.Body.String(), message) { t.Fatal("unsupported placement attempted credential decryption", response.Body) } - assertSelfHostedMCPRejectionHasNoWrites(t, db.pool, tenant) + assertSelfHostedMCPRejectionHasNoWrites(t, s.pool, tenant) } }) } diff --git a/services/core/tests/integration/remote_mcp_test.go b/services/core/tests/integration/remote_mcp_test.go index dec32b273..0344b0c0a 100644 --- a/services/core/tests/integration/remote_mcp_test.go +++ b/services/core/tests/integration/remote_mcp_test.go @@ -15,8 +15,8 @@ import ( ) func TestSelfHostedServiceMCPRejectedWithoutWrites(t *testing.T) { - s, db, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) - handler := selfHostedMCPAdmissionHandler(t, s, db, tenant) + s, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) + handler := selfHostedMCPAdmissionHandler(t, s, tenant) for _, mode := range []string{"unattached", "missing", "wrong URL", "foreign Vault", "anonymous", "implicit", "explicit", "required anonymous", "required bearer"} { for _, initial := range []bool{false, true} { @@ -67,20 +67,20 @@ func TestSelfHostedServiceMCPRejectedWithoutWrites(t *testing.T) { t.Fatal("rejected request exposed private authentication") } - assertSelfHostedMCPRejectionHasNoWrites(t, db.pool, tenant) + assertSelfHostedMCPRejectionHasNoWrites(t, s.pool, tenant) } } } -func selfHostedMCPAdmissionFixture(t *testing.T) (*Store, fixtureDB, string, vaults.Vault, vaults.Credential) { +func selfHostedMCPAdmissionFixture(t *testing.T) (*Store, string, vaults.Vault, vaults.Credential) { t.Helper() _, pool := testStore(t) cipher, err := credentialcrypto.New([]byte(strings.Repeat("k", 32))) if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} - _, service, err := fixtureVaults(db) + s := NewWithCredentialCipher(pool, cipher) + _, service, err := fixtureVaults(s) if err != nil { t.Fatal(err) } @@ -93,13 +93,13 @@ func selfHostedMCPAdmissionFixture(t *testing.T) (*Store, fixtureDB, string, vau if err != nil { t.Fatal(err) } - return s, db, tenant, vault, credential + return s, tenant, vault, credential } -func selfHostedMCPAdmissionHandler(t *testing.T, s *Store, db fixtureDB, tenant string) http.Handler { +func selfHostedMCPAdmissionHandler(t *testing.T, s *Store, tenant string) http.Handler { t.Helper() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "test", TenantID: tenant, TokenSHA256: runtimedevice.HashCredential("test-token")}}) - handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) + handler, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/request_body_public_test.go b/services/core/tests/integration/request_body_public_test.go index 37073f5a3..f075d4ce3 100644 --- a/services/core/tests/integration/request_body_public_test.go +++ b/services/core/tests/integration/request_body_public_test.go @@ -29,7 +29,7 @@ func TestRequestBodyGateRejectsWithoutWritesPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) owner, foreign, ownerTenant := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "body-owner", TokenSHA256: runtimedevice.HashCredential(owner), TenantID: ownerTenant}, @@ -37,7 +37,7 @@ func TestRequestBodyGateRejectsWithoutWritesPostgres(t *testing.T) { }) // No Runtime is connected, so a file write that passes the gate is unavailable. unavailable := func(d *api.Dependencies) { d.Execution.Workspaces = unavailableWorkspaces{strictStandIn{t}} } - h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), unavailable, acceptUnavailable(t)) + h, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), unavailable, acceptUnavailable(t)) if err != nil { t.Fatal(err) } @@ -170,10 +170,10 @@ func TestRequestBodyGateExcludedRoutesPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) token, tenant := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "excluded-owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}}) - 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) } diff --git a/services/core/tests/integration/root_fixture_test.go b/services/core/tests/integration/root_fixture_test.go index 1696dcdf1..2fde53550 100644 --- a/services/core/tests/integration/root_fixture_test.go +++ b/services/core/tests/integration/root_fixture_test.go @@ -12,6 +12,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -29,9 +30,14 @@ type Store struct { placement *placement.Rules } +// defaultPlacement is the placement rules cmd/server builds on the built-in +// providers and an unset public URL. +var defaultPlacement, _ = placement.NewRules(providers.Builtin(), "") + +// New is the fixture on pool without a credential key, under defaultPlacement. func New(pool *pgxpool.Pool) *Store { pooled := pgunit.NewPool(pool) - return &Store{queries: sqlc.New(pool), pool: pool, pooled: pooled, writer: pooled} + return &Store{queries: sqlc.New(pool), pool: pool, pooled: pooled, writer: pooled, placement: defaultPlacement} } // NewWithCredentialCipher is New with a credential key, so Sessions can freeze diff --git a/services/core/tests/integration/runtime_capabilities_pending_test.go b/services/core/tests/integration/runtime_capabilities_pending_test.go index f0c854b16..5bbe67cf6 100644 --- a/services/core/tests/integration/runtime_capabilities_pending_test.go +++ b/services/core/tests/integration/runtime_capabilities_pending_test.go @@ -20,8 +20,7 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(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) tenant := uuid.NewString() session, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{ Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), @@ -31,15 +30,15 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { if err != nil { t.Fatal(err) } - env, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) + env, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } provider := &initializingProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, initializationPeer: initializationPeer{deferred: true}} key := uuid.NewString() - worker, _ := managedWorkerMode(t, s, db, key, provider, false, true) + worker, _ := managedWorkerMode(t, s, key, provider, false, true) owner, err := worker.ProvisionEnvironment(t.Context(), tenant, env.ID, key) - if err != nil || initializationState(t, db.pool, owner.TenantID, owner.EnvironmentID) != "pending" { + if err != nil || initializationState(t, s, owner.TenantID, owner.EnvironmentID) != "pending" { t.Fatal(owner, err) } for range 4 { @@ -47,8 +46,8 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { t.Fatal(err) } } - owner, err = fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) - if err != nil || initializationState(t, db.pool, owner.TenantID, owner.EnvironmentID) != "pending" || owner.State != "running" || provider.writes.Load() != 0 || provider.kills != 0 { + owner, err = deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + if err != nil || initializationState(t, s, owner.TenantID, owner.EnvironmentID) != "pending" || owner.State != "running" || provider.writes.Load() != 0 || provider.kills != 0 { t.Fatal("missing socket consumed initialization or requested cleanup", owner, err, provider.writes.Load(), provider.kills) } if _, err := sessionAdapter(s).GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { diff --git a/services/core/tests/integration/runtime_compute_lifecycle_test.go b/services/core/tests/integration/runtime_compute_lifecycle_test.go index 01700b449..7d3286718 100644 --- a/services/core/tests/integration/runtime_compute_lifecycle_test.go +++ b/services/core/tests/integration/runtime_compute_lifecycle_test.go @@ -260,7 +260,6 @@ func (p *fakeCheckpointProvider) connect(ctx context.Context, b sandbox.Bootstra type computeLifecycleFixture struct { t *testing.T store *Store - db fixtureDB provider *fakeCheckpointProvider worker *execution.Worker stop func() @@ -270,10 +269,10 @@ type computeLifecycleFixture struct { func newComputeLifecycleFixture(t *testing.T, maxActive, maxRetained int) *computeLifecycleFixture { t.Helper() - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) registry := runtimegateway.NewRegistry() p := &fakeCheckpointProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, computes: map[string]sandbox.ComputeState{}, snapshots: map[string]sandbox.SnapshotIdentity{}, bootstraps: map[string]sandbox.Bootstrap{}, peers: map[string]*websocket.Conn{}, registry: registry} - handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(fixtureSessionStore(db)), Registry: registry}) + handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(sessionAdapter(s)), Registry: registry}) server := httptest.NewServer(http.HandlerFunc(handler.WS)) p.endpoint = "ws" + strings.TrimPrefix(server.URL, "http") t.Cleanup(func() { @@ -284,7 +283,7 @@ func newComputeLifecycleFixture(t *testing.T, maxActive, maxRetained int) *compu } server.Close() }) - f := &computeLifecycleFixture{t: t, store: s, db: db, provider: p, key: uuid.NewString(), policy: execution.RuntimeSuspensionPolicy{IdleTimeout: time.Second, Retention: time.Hour, MaxActive: maxActive, MaxRetained: maxRetained}} + f := &computeLifecycleFixture{t: t, store: s, provider: p, key: uuid.NewString(), policy: execution.RuntimeSuspensionPolicy{IdleTimeout: time.Second, Retention: time.Hour, MaxActive: maxActive, MaxRetained: maxRetained}} f.start() return f } @@ -297,7 +296,7 @@ func (f *computeLifecycleFixture) start() { var w *execution.Worker var err error for { - w, err = startWorkerErr(t.Context(), f.db, dispatcher) + w, err = startWorkerErr(t.Context(), f.store, dispatcher) if err == nil || !errors.Is(err, pgunit.ErrLeaseHeld) || !time.Now().Before(deadline) { break } @@ -315,14 +314,14 @@ func (f *computeLifecycleFixture) start() { } func (f *computeLifecycleFixture) sql(query string, args ...any) { f.t.Helper() - if _, err := f.db.pool.Exec(f.t.Context(), query, args...); err != nil { + if _, err := f.store.pool.Exec(f.t.Context(), query, args...); err != nil { f.t.Fatal(err) } } func (f *computeLifecycleFixture) create() (string, sessions.Session, sessions.Environment, deployment.Allocation) { t := f.t t.Helper() - tenant, session, environment := managedSession(t, f.store, f.db) + tenant, session, environment := managedSession(t, f.store) owner, err := f.worker.ProvisionEnvironment(t.Context(), tenant, environment.ID, f.key) if err != nil { t.Fatal(err) @@ -346,7 +345,7 @@ func (f *computeLifecycleFixture) phase(tenant, environment, phase string) deplo if err != nil { f.t.Fatal(err) } - owner, err = fixtureReader(f.db).EnvironmentAllocation(f.t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) + owner, err = deploymentStore(f.store).EnvironmentAllocation(f.t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) if err != nil { f.t.Fatal(err) } @@ -389,7 +388,7 @@ func TestRuntimeComputeLifecycleIdleSuspendAndQueuedSameSessionWake(t *testing.T t.Fatal("wake replaced Session or replayed allocation") } var completedCount, queuedCount int - if err := f.db.pool.QueryRow(t.Context(), `SELECT count(*) FILTER(WHERE id=$2 AND status='completed'),count(*) FILTER(WHERE id=$3 AND status='queued') FROM turns WHERE session_id=$1`, session.ID, completed, queued).Scan(&completedCount, &queuedCount); err != nil || completedCount != 1 || queuedCount != 1 { + if err := f.store.pool.QueryRow(t.Context(), `SELECT count(*) FILTER(WHERE id=$2 AND status='completed'),count(*) FILTER(WHERE id=$3 AND status='queued') FROM turns WHERE session_id=$1`, session.ID, completed, queued).Scan(&completedCount, &queuedCount); err != nil || completedCount != 1 || queuedCount != 1 { t.Fatal("wake replayed/consumed prior or next Turn", err) } if got, err := sessionAdapter(f.store).GetSession(t.Context(), tenant, session.ID); err != nil || string(got.Configuration) != string(session.Configuration) { @@ -454,7 +453,7 @@ func TestRuntimeComputeLifecycleQuiesceRejectionAndUnknownIntentRollback(t *test } // A restart with persisted quiescing but no acknowledgement must resume the // source. Inject only the durable phase, never a fake snapshot or new source. - current, err := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + current, err := deploymentStore(f.store).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil { t.Fatal(err) } @@ -476,7 +475,7 @@ func TestRuntimeComputeLifecycleWakeDuringQuiesceResumesSource(t *testing.T) { tenant, _, env, owner := f.create() f.complete(owner) wakeResult := make(chan error, 1) - activity := fixtureDeployment(t, f.db) + activity := deploymentService(t, f.store) f.provider.mu.Lock() f.provider.beforeQuiesce = func() { wakeResult <- activity.TouchActivity(t.Context(), tenant, env.ID) @@ -515,11 +514,11 @@ func TestRuntimeComputeLifecycleSuspendedDeletionAndExpiryCleanup(t *testing.T) } else { f.sql(`UPDATE runtime_allocations SET compute_retained_until=clock_timestamp()-interval '1 second' WHERE id=$1`, owner.ID) } - reconcileManagedState(t, f.worker, f.db, tenant, env.ID, "released") + reconcileManagedState(t, f.worker, f.store, tenant, env.ID, "released") if len(f.provider.computes) != 0 || len(f.provider.snapshots) != 0 || f.provider.snapshotDeletes != 1 { t.Fatal("retained snapshot survived cleanup") } - if _, ok, err := fixtureSessionStore(f.db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := sessionAdapter(f.store).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("cleanup retained daemon authority", err) } }) @@ -529,7 +528,7 @@ func TestRuntimeComputeLifecycleSuspendedDeletionAndExpiryCleanup(t *testing.T) func TestRuntimeComputeLifecycleCapacityBoundsActiveAndRetained(t *testing.T) { f := newComputeLifecycleFixture(t, 1, 2) tenant, _, env, owner := f.create() - tenant2, _, env2 := managedSession(t, f.store, f.db) + tenant2, _, env2 := managedSession(t, f.store) if _, err := f.worker.ProvisionEnvironment(t.Context(), tenant2, env2.ID, f.key); !errors.Is(err, execution.ErrExecutionUnavailable) { t.Fatalf("active capacity ignored: %v", err) } @@ -546,7 +545,7 @@ func TestRuntimeComputeLifecycleCapacityBoundsActiveAndRetained(t *testing.T) { for range 4 { f.worker.ReconcileManagedRuntimes(t.Context()) } - first, err := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + first, err := deploymentStore(f.store).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || first.ComputePhase != "suspended" || f.provider.restores != 0 { t.Fatal("wake exceeded active capacity", err) } @@ -561,7 +560,7 @@ func TestRuntimeComputeLifecycleCapacityBoundsActiveAndRetained(t *testing.T) { f.complete(second) f.phase(tenant2, env2.ID, "suspended") // Both retained allocations count even when their source VMs are gone. - tenant3, _, env3 := managedSession(t, f.store, f.db) + tenant3, _, env3 := managedSession(t, f.store) if _, err := f.worker.ProvisionEnvironment(t.Context(), tenant3, env3.ID, f.key); !errors.Is(err, execution.ErrExecutionUnavailable) { t.Fatalf("retained capacity ignored: %v", err) } diff --git a/services/core/tests/integration/runtime_configuration_cleanup_test.go b/services/core/tests/integration/runtime_configuration_cleanup_test.go index 3ea9cb028..2196f7674 100644 --- a/services/core/tests/integration/runtime_configuration_cleanup_test.go +++ b/services/core/tests/integration/runtime_configuration_cleanup_test.go @@ -83,7 +83,7 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { {name: "kill unavailable stays retained", inspectionError: sandbox.ErrInvalid, killError: sandbox.ErrComputeUnconfirmed, wantSettled: true, wantKill: true}, } { t.Run(test.name, func(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) key := uuid.NewString() p := &configurationCleanupProvider{ lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}, loseCreate: test.loseCreate}, @@ -91,8 +91,8 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { settleInspection: test.settleInspection, foreign: test.foreign, inspectionError: test.inspectionError, killError: test.killError, } - w, _ := managedWorker(t, s, db, key, p) - tenant, session, environment := managedSession(t, s, db) + w, _ := managedWorker(t, s, key, p) + tenant, session, environment := managedSession(t, s) owner, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key) if (err != nil) != (test.rejectCreate || test.loseCreate) || owner.ID == "" { t.Fatal("unexpected creation outcome", owner, err) @@ -108,7 +108,7 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - before, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) + before, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || before.State == "cleanup_pending" || before.State == "released" { t.Fatal("drift authorized cleanup", before, err) } @@ -129,17 +129,17 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { if test.wantReleased { wantState = "released" } - reconcileManagedState(t, w, db, tenant, environment.ID, wantState) + reconcileManagedState(t, w, s, tenant, environment.ID, wantState) // A second observation must not turn absence after Kill into proof // that an unknown original Create can no longer mutate resources. if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) + got, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || got.ID != owner.ID || got.State != wantState || got.CreateSettled != test.wantSettled { t.Fatal("cleanup lost ownership or settlement", got, err) } - if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("cleanup retained execution authority", err) } p.mu.Lock() diff --git a/services/core/tests/integration/runtime_connection_test.go b/services/core/tests/integration/runtime_connection_test.go index 85d5ccd11..19947a340 100644 --- a/services/core/tests/integration/runtime_connection_test.go +++ b/services/core/tests/integration/runtime_connection_test.go @@ -21,11 +21,11 @@ import ( ) func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { - s, db := newManagedTestStoreDB(t) - tenant, session, environment := managedSession(t, s, db) + s, _ := newManagedTestStore(t) + tenant, session, environment := managedSession(t, s) 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) } @@ -35,7 +35,7 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { p := &lifecycleProvider{resources: map[string]sandbox.Info{}} key := uuid.NewString() start := func() *execution.Worker { - w := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: server.URL + "/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p}}) + w := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: server.URL + "/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p}}) return w } stop := func(w *execution.Worker) { @@ -55,7 +55,7 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - got, err := fixtureSessionStore(db).GetEnvironment(t.Context(), tenant, environment.ID) + got, err := sessionAdapter(s).GetEnvironment(t.Context(), tenant, environment.ID) if err != nil { t.Fatal(err) } @@ -91,7 +91,7 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { if err != nil || got.LastTurn != nil || got.EnvironmentInputActivity != nil { t.Fatal("connection fabricated native execution", err) } - if _, err := fixtureSessionStore(db).GetEnvironment(t.Context(), uuid.NewString(), environment.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetEnvironment(t.Context(), uuid.NewString(), environment.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign Environment access", err) } p.unavailable = true @@ -107,14 +107,14 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { stop(w) w = start() assertStatus("connected") - retained, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) + retained, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || retained.ID != owner.ID || retained.DeviceID != owner.DeviceID || p.creates != 1 { t.Fatal("restart replaced Runtime identity", err) } if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } - reconcileManagedState(t, w, db, tenant, environment.ID, "released") + reconcileManagedState(t, w, s, tenant, environment.ID, "released") if conn, err := dial(p.credential); err == nil { conn.Close() t.Fatal("released Runtime reconnected") diff --git a/services/core/tests/integration/runtime_creation_settlement_test.go b/services/core/tests/integration/runtime_creation_settlement_test.go index ea035f854..b827e9c49 100644 --- a/services/core/tests/integration/runtime_creation_settlement_test.go +++ b/services/core/tests/integration/runtime_creation_settlement_test.go @@ -41,11 +41,11 @@ func (p *absentCreationProvider) GetInfo(_ context.Context, r sandbox.Reference) func TestManagedRuntimeConfirmedAbsentCreateReleasesAtomically(t *testing.T) { for _, cancelled := range []bool{false, true} { t.Run(map[bool]string{false: "live caller", true: "cancelled caller"}[cancelled], func(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) key := uuid.NewString() p := &absentCreationProvider{} - w, _ := managedWorker(t, s, db, key, p) - tenant, session, environment := managedSession(t, s, db) + w, _ := managedWorker(t, s, key, p) + tenant, session, environment := managedSession(t, s) ctx, cancel := context.WithCancel(t.Context()) defer cancel() if cancelled { @@ -58,11 +58,11 @@ func TestManagedRuntimeConfirmedAbsentCreateReleasesAtomically(t *testing.T) { if p.kills != 0 || p.creates != 1 { t.Fatal("absence proof still called external cleanup", p.kills, p.creates) } - stored, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) + stored, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || stored.State != "released" || !stored.CreateSettled { t.Fatal("release not durable", err) } - if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("released credential retained authority", err) } value, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) @@ -79,32 +79,32 @@ func TestManagedRuntimeConfirmedAbsentCreateReleasesAtomically(t *testing.T) { } } func TestManagedRuntimeForeignAbsenceCannotReleaseCreation(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) key := uuid.NewString() p := &absentCreationProvider{foreign: true} - w, _ := managedWorker(t, s, db, key, p) - tenant, _, environment := managedSession(t, s, db) + w, _ := managedWorker(t, s, key, p) + tenant, _, environment := managedSession(t, s) if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err == nil { t.Fatal("foreign absence accepted") } - owner, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) + owner, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || owner.CreateSettled || owner.State == "released" { t.Fatal("foreign proof settled original attempt", owner, err) } } func TestManagedRuntimeObservedSettlementAllowsOwnedCleanup(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) key := uuid.NewString() p := &absentCreationProvider{observeSettled: true} - w, _ := managedWorker(t, s, db, key, p) - tenant, session, environment := managedSession(t, s, db) + w, _ := managedWorker(t, s, key, p) + tenant, session, environment := managedSession(t, s) if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err == nil { t.Fatal("uncertain Create succeeded") } if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } - reconcileManagedState(t, w, db, tenant, environment.ID, "released") + reconcileManagedState(t, w, s, tenant, environment.ID, "released") if p.kills != 1 || p.creates != 1 { t.Fatal("settled observation replayed Create or skipped cleanup", p.creates, p.kills) } diff --git a/services/core/tests/integration/runtime_deployment_worker_test.go b/services/core/tests/integration/runtime_deployment_worker_test.go index ab326a2f4..42ce2fae8 100644 --- a/services/core/tests/integration/runtime_deployment_worker_test.go +++ b/services/core/tests/integration/runtime_deployment_worker_test.go @@ -12,11 +12,11 @@ import ( ) func TestManagedDeploymentStartupRejectsSwitchBeforeBackendAccess(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) key := uuid.NewString() old := &lifecycleProvider{resources: map[string]sandbox.Info{}} - worker, stop := managedWorker(t, s, db, key, old) - tenant, _, environment := managedSession(t, s, db) + worker, stop := managedWorker(t, s, key, old) + tenant, _, environment := managedSession(t, s) owner, err := worker.ProvisionEnvironment(t.Context(), tenant, environment.ID, key) if err != nil { t.Fatal(err) @@ -25,13 +25,13 @@ func TestManagedDeploymentStartupRejectsSwitchBeforeBackendAccess(t *testing.T) replacement := &lifecycleProvider{resources: map[string]sandbox.Info{}} config := &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: uuid.NewString(), BackendFingerprint: strings.Repeat("b", 64), Provider: replacement, AdmissionPaused: true} start := func(config *execution.RuntimeProvider) error { - _, err := startWorkerErr(t.Context(), db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: config}) + _, err := startWorkerErr(t.Context(), s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: config}) return err } if err := start(config); err == nil || !strings.Contains(err.Error(), "maintenance") { t.Fatal("startup switched active deployment", err) } - worker, stop = managedWorkerMode(t, s, db, key, old, true) + worker, stop = managedWorkerMode(t, s, key, old, true) replay, err := worker.ProvisionEnvironment(t.Context(), tenant, environment.ID, key) if err != nil || !replay.Replayed || replay.ID != owner.ID { t.Fatal("maintenance interrupted existing allocation", replay, err) @@ -46,11 +46,11 @@ func TestManagedDeploymentStartupRejectsSwitchBeforeBackendAccess(t *testing.T) if replacement.creates != 0 || replacement.kills != 0 { t.Fatal("rejected startup touched new backend") } - got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) + got, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || got.ID != owner.ID || got.ProviderKey != key { t.Fatal("rejected startup rewrote resource owner", got, err) } // Failed startup relinquishes its lease, so the original backend can resume. - _, stop = managedWorker(t, s, db, key, old) + _, stop = managedWorker(t, s, key, old) stop() } diff --git a/services/core/tests/integration/runtime_enrollment_connection_test.go b/services/core/tests/integration/runtime_enrollment_connection_test.go index 9a452f384..791fe6dd9 100644 --- a/services/core/tests/integration/runtime_enrollment_connection_test.go +++ b/services/core/tests/integration/runtime_enrollment_connection_test.go @@ -22,7 +22,7 @@ import ( ) func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { - s, db := newTestStoreDB(t) + s, _ := testStore(t) principal := FixtureExecutorPrincipal(t, s, uuid.NewString()) session, err := s.CreateSession(t.Context(), principal.TenantID, sessions.CreateSession{ Creator: principal.Subject(), Engine: "codex", IdempotencyKey: uuid.NewString(), @@ -31,26 +31,26 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) if err != nil { t.Fatal(err) } - key, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + key, err := sessionService(t, s).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } digest := sha256.Sum256([]byte(key.Token)) - bound, err := fixtureSessionService(t, db).EnrollRuntime(t.Context(), environment.ID, hex.EncodeToString(digest[:])) + bound, err := sessionService(t, s).EnrollRuntime(t.Context(), environment.ID, hex.EncodeToString(digest[:])) if err != nil { t.Fatal(err) } 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) } - connection := runtimeenrollment.ConnectionHandler(fixtureSessionStore(db), registry) + connection := runtimeenrollment.ConnectionHandler(sessionAdapter(s), registry) assertConnection := func(target, token, status string, code int) { t.Helper() request := httptest.NewRequest("GET", "/api/v1/agent-daemon/connection?environment_id="+target, nil) @@ -69,13 +69,13 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { } assertConnection(environment.ID, key.Token, "disconnected", 200) assertConnection(uuid.NewString(), key.Token, "", 401) - otherKey, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + otherKey, err := sessionService(t, s).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } assertConnection(environment.ID, otherKey.Token, "", 409) foreign := FixtureExecutorPrincipal(t, s, uuid.NewString()) - foreignKey, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), foreign, uuid.NewString(), "") + foreignKey, err := sessionService(t, s).IssueExecutorCredential(t.Context(), foreign, uuid.NewString(), "") if err != nil { t.Fatal(err) } @@ -84,7 +84,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { server.Start() t.Cleanup(func() { server.Close(); runtime.CloseConnections(registry) }) start := func() func() { - worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: registry}) + worker := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: registry}) ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -120,7 +120,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { t.Helper() deadline := time.Now().Add(5 * time.Second) for time.Now().Before(deadline) { - current, err := fixtureSessionStore(db).GetEnvironment(t.Context(), principal.TenantID, environment.ID) + current, err := sessionAdapter(s).GetEnvironment(t.Context(), principal.TenantID, environment.ID) if err == nil && current.Status == status { return } @@ -131,7 +131,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { first := connect(key.Token) await("connected") assertConnection(environment.ID, key.Token, "connected", 200) - rotated, err := fixtureSessionService(t, db).RotateExecutorCredential(t.Context(), principal, key.KeyID) + rotated, err := sessionService(t, s).RotateExecutorCredential(t.Context(), principal, key.KeyID) if err != nil { t.Fatal(err) } @@ -146,7 +146,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { second := connect(rotated.Token) await("connected") assertConnection(environment.ID, rotated.Token, "connected", 200) - awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, db.pool) + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, s.pool) stop() stop = nil awaitRelease() @@ -154,7 +154,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { // live, authorized daemon. No compute allocation or native execution is made. stop = start() await("connected") - if err = fixtureSessionService(t, db).RevokeExecutorCredential(t.Context(), principal, key.KeyID); err != nil { + if err = sessionService(t, s).RevokeExecutorCredential(t.Context(), principal, key.KeyID); err != nil { t.Fatal(err) } assertConnection(environment.ID, rotated.Token, "", 401) @@ -164,7 +164,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { t.Fatal("revoked socket retained authority") } var allocations int - if err = db.pool.QueryRow(t.Context(), "SELECT count(*) FROM runtime_allocations WHERE environment_id=$1", environment.ID).Scan(&allocations); err != nil || allocations != 0 { + if err = s.pool.QueryRow(t.Context(), "SELECT count(*) FROM runtime_allocations WHERE environment_id=$1", environment.ID).Scan(&allocations); err != nil || allocations != 0 { t.Fatal("user Runtime acquired managed allocation", allocations, err) } current, err := sessionAdapter(s).GetSession(t.Context(), principal.TenantID, session.ID) diff --git a/services/core/tests/integration/runtime_initialization_test.go b/services/core/tests/integration/runtime_initialization_test.go index ba45c2311..57a4bf5cc 100644 --- a/services/core/tests/integration/runtime_initialization_test.go +++ b/services/core/tests/integration/runtime_initialization_test.go @@ -51,8 +51,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(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) tenant := uuid.NewString() input := sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"environment":{"type":"openai_hosted"}}`), InitialFiles: []environmentconfig.InitialFile{{Type: "inline", Path: "/workspace/a", Data: []byte("first")}, {Type: "inline", Path: "/workspace/b", Data: []byte("second")}}} if setupOnly { @@ -64,7 +63,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { if err != nil { t.Fatal(err) } - env, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) + env, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -75,7 +74,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { } p := &initializingProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, initializationPeer: initializationPeer{deferred: true}} p.apply = func(_ proto.RuntimePreparePayload, _ []byte) proto.RuntimePrepareResultPayload { - if _, err := fixtureSessionStore(db).GetSessionDevice(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSessionDevice(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Error("premature file access", err) } if _, err := sessionAdapter(s).GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { @@ -87,9 +86,9 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { return completedInitialization(proto.RuntimePreparePayload{}, nil) } key := uuid.NewString() - w, stop := managedWorkerMode(t, s, db, key, p, false, true) + w, stop := managedWorkerMode(t, s, key, p, false, true) if mode == "restart" { - awaitInitialization(t, db.pool, tenant, env.ID, "failed") + awaitInitialization(t, s, tenant, env.ID, "failed") if p.writes.Load() != 0 { t.Fatal("recovered unknown operation replayed") } @@ -99,11 +98,11 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { if err != nil { t.Fatal(err) } - if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || !ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || !ok { t.Fatal("preparation blocked authentication", err) } time.Sleep(350 * time.Millisecond) - if initializationState(t, db.pool, tenant, env.ID) != "pending" || p.writes.Load() != 0 { + if initializationState(t, s, tenant, env.ID) != "pending" || p.writes.Load() != 0 { t.Fatal("missing socket consumed initialization") } p.deferred = false @@ -114,7 +113,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { if mode == "uncertain" { want = "failed" } - awaitInitialization(t, db.pool, tenant, env.ID, want) + awaitInitialization(t, s, tenant, env.ID, want) if mode == "complete" { if int(p.writes.Load()) != expectedSteps { t.Fatal("missing operations", p.writes.Load()) @@ -123,7 +122,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { t.Fatal("completed preparation blocked", err) } stop() - _, _ = managedWorkerMode(t, s, db, key, p, false, true) + _, _ = managedWorkerMode(t, s, key, p, false, true) time.Sleep(350 * time.Millisecond) if int(p.writes.Load()) != expectedSteps { t.Fatal("completed preparation replayed") @@ -142,7 +141,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { } func TestManagedRuntimePreparationAllOperationsUsePeer(t *testing.T) { - s, db := hostedFailureStore(t) + s := hostedFailureStore(t) var archive bytes.Buffer writer := zip.NewWriter(&archive) for path, body := range map[string]string{"proof/.codex-plugin/plugin.json": `{"name":"plugin","description":"A plugin.","skills":"./skills"}`, "proof/skills/example/SKILL.md": "---\nname: plugin-proof\ndescription: A plugin Skill.\n---\nProof."} { @@ -159,7 +158,7 @@ func TestManagedRuntimePreparationAllOperationsUsePeer(t *testing.T) { } tenant := uuid.NewString() fileBody := bytes.Repeat([]byte("bounded bytes"), 12000) - session, environment := hostedFailureSession(t, s, db, tenant, sessions.CreateSession{ + session, environment := hostedFailureSession(t, s, tenant, sessions.CreateSession{ InitialFiles: []environmentconfig.InitialFile{{Type: "inline", Path: "/workspace/first", Data: fileBody}}, Initialization: environmentconfig.Setup{Skills: []environmentconfig.Skill{hostedFailureSkill(t)}, Plugins: []environmentconfig.Plugin{{Metadata: agentplugin.Metadata{Type: "inline", Name: "plugin", Description: "A plugin."}, Archive: archive.Bytes()}}, Packages: v1.EnvironmentPackages{NPM: []string{"is-number@7.0.0"}, Python: []string{"packaging==24.2"}}, Commands: []environmentconfig.SetupCommand{{Command: "read installed bundles and create directory"}}, CapabilityDirectories: []string{"/workspace/generated"}}, }) @@ -183,19 +182,19 @@ func TestManagedRuntimePreparationAllOperationsUsePeer(t *testing.T) { return completedInitialization(request, data) } key := uuid.NewString() - worker, _ := managedWorkerMode(t, s, db, key, provider, false, true) + worker, _ := managedWorkerMode(t, s, key, provider, false, true) if _, err := worker.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err != nil { t.Fatal(err) } - awaitInitialization(t, db.pool, tenant, environment.ID, "complete") + awaitInitialization(t, s, tenant, environment.ID, "complete") actionsMu.Lock() defer actionsMu.Unlock() expected := []string{"file", "configure", "skill", "plugin", "npm", "python", "setup", "finalize"} if !reflect.DeepEqual(actions, expected) || provider.commandCalls.Load() != 0 { t.Fatal("typed ordering or provider isolation", actions, provider.commandCalls.Load()) } - allocation, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) - if err != nil || initializationState(t, db.pool, allocation.TenantID, allocation.EnvironmentID) != "complete" { + allocation, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) + if err != nil || initializationState(t, s, allocation.TenantID, allocation.EnvironmentID) != "complete" { t.Fatal("initialization incomplete", allocation, err) } } diff --git a/services/core/tests/integration/runtime_input_admission_test.go b/services/core/tests/integration/runtime_input_admission_test.go index 6ff20dd1d..66617e3db 100644 --- a/services/core/tests/integration/runtime_input_admission_test.go +++ b/services/core/tests/integration/runtime_input_admission_test.go @@ -13,15 +13,15 @@ import ( ) func TestManagedRuntimeMaintenancePreservesCancelAndRetry(t *testing.T) { - s, db := newManagedTestStoreDB(t) - tenant, session, _ := managedSession(t, s, db) + s, _ := newManagedTestStore(t) + tenant, session, _ := managedSession(t, s) inputs := []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"accepted work"}`)}} accepted, err := submitInputs(t.Context(), s, tenant, session.ID, "work", inputs) if err != nil { t.Fatal(err) } p := &lifecycleProvider{resources: map[string]sandbox.Info{}} - w, stop := managedWorkerMode(t, s, db, uuid.NewString(), p, true) + w, stop := managedWorkerMode(t, s, uuid.NewString(), p, true) defer stop() if _, err := w.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: session.Configuration}); !errors.Is(err, placement.ErrAdmissionClosed) { t.Fatal("maintenance accepted new hosted Session", err) diff --git a/services/core/tests/integration/runtime_lifecycle_test.go b/services/core/tests/integration/runtime_lifecycle_test.go index 67ed0160d..f03f85e0a 100644 --- a/services/core/tests/integration/runtime_lifecycle_test.go +++ b/services/core/tests/integration/runtime_lifecycle_test.go @@ -72,23 +72,23 @@ func (p *lifecycleProvider) RunCommand(context.Context, sandbox.Reference, sandb return sandbox.CommandResult{}, errors.New("not used") } -func managedWorker(t *testing.T, s *Store, db fixtureDB, key string, p sandbox.SandboxProvider) (*execution.Worker, func()) { +func managedWorker(t *testing.T, s *Store, key string, p sandbox.SandboxProvider) (*execution.Worker, func()) { t.Helper() - return managedWorkerMode(t, s, db, key, p, false) + return managedWorkerMode(t, s, key, p, false) } -func managedWorkerMode(t *testing.T, s *Store, db fixtureDB, key string, p sandbox.SandboxProvider, maintenance bool, run ...bool) (*execution.Worker, func()) { +func managedWorkerMode(t *testing.T, s *Store, key string, p sandbox.SandboxProvider, maintenance bool, run ...bool) (*execution.Worker, func()) { t.Helper() registry := runtimegateway.NewRegistry() if peer, ok := p.(interface { setRuntimeGateway(*testing.T, string, *runtimegateway.Registry) }); ok { - handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(fixtureSessionStore(db)), Registry: registry}) + handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(sessionAdapter(s)), Registry: registry}) server := httptest.NewServer(http.HandlerFunc(handler.WS)) t.Cleanup(server.Close) peer.setRuntimeGateway(t, "ws"+strings.TrimPrefix(server.URL, "http"), registry) } - w := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p, AdmissionPaused: maintenance}}) + w := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p, AdmissionPaused: maintenance}}) if len(run) > 0 && run[0] { ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) @@ -113,27 +113,27 @@ func managedWorkerMode(t *testing.T, s *Store, db fixtureDB, key string, p sandb return w, stop } -func managedSession(t *testing.T, s *Store, db fixtureDB) (string, sessions.Session, sessions.Environment) { +func managedSession(t *testing.T, s *Store) (string, sessions.Session, sessions.Environment) { t.Helper() tenant := uuid.NewString() v, e := s.CreateSession(t.Context(), tenant, WithFixtureModelProvider(sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"model":"test"},"environment":{"type":"openai_hosted","network":{"access":"enabled"}}}`)})) if e != nil { t.Fatal(e) } - env, e := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, v.ID) + env, e := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, v.ID) if e != nil { t.Fatal(e) } return tenant, v, env } -func reconcileManagedState(t *testing.T, w *execution.Worker, db fixtureDB, tenant, environment, state string) { +func reconcileManagedState(t *testing.T, w *execution.Worker, s *Store, tenant, environment, state string) { t.Helper() for range 100 { if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) + got, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) if err != nil { t.Fatal(err) } @@ -145,23 +145,23 @@ func reconcileManagedState(t *testing.T, w *execution.Worker, db fixtureDB, tena } func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { - s, db := newManagedTestStoreDB(t) - tenant, session, env := managedSession(t, s, db) + s, _ := newManagedTestStore(t) + tenant, session, env := managedSession(t, s) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}, loseCreate: true} - w, stop := managedWorker(t, s, db, key, p) + w, stop := managedWorker(t, s, key, p) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err == nil || owner.ID == "" { t.Fatal("fault did not retain allocation") } - credential, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID) + credential, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID) if err != nil || !ok || credential.CredentialHash != p.credentialHash { t.Fatal("provider received unbound credential") } stop() - next, _ := managedWorker(t, s, db, key, p) - reconcileManagedState(t, next, db, tenant, env.ID, "running") - recovered, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + next, _ := managedWorker(t, s, key, p) + reconcileManagedState(t, next, s, tenant, env.ID, "running") + recovered, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || recovered.ID != owner.ID || !recovered.CreateSettled || recovered.State != "running" { t.Fatalf("lost response recovery: %+v %v", recovered, err) } @@ -173,22 +173,22 @@ func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { t.Fatal(err) } // A scan may first exhaust its previous cursor before starting a new cycle. - reconcileManagedState(t, next, db, tenant, env.ID, "released") - clean, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + reconcileManagedState(t, next, s, tenant, env.ID, "released") + clean, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || clean.State != "released" || p.kills != 1 { t.Fatalf("deleted cleanup: %+v %v", clean, err) } - if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("cleanup did not revoke authority") } } func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { - s, db := newManagedTestStoreDB(t) - tenant, session, env := managedSession(t, s, db) + s, _ := newManagedTestStore(t) + tenant, session, env := managedSession(t, s) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}, loseCreate: true, absent: true} - w, _ := managedWorker(t, s, db, key, p) + w, _ := managedWorker(t, s, key, p) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err == nil { t.Fatal("expected uncertain creation") @@ -196,43 +196,43 @@ func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } - reconcileManagedState(t, w, db, tenant, env.ID, "cleanup_pending") - got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + reconcileManagedState(t, w, s, tenant, env.ID, "cleanup_pending") + got, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "cleanup_pending" || got.CreateSettled || p.creates != 1 { t.Fatalf("unknown creation forgotten: %+v %v", got, err) } - if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("unknown allocation retains execution authority") } // A late completion is still owned and reclaimed on the next scan. p.resources[owner.ID] = sandbox.Info{Reference: sandbox.Reference{TenantID: tenant, EnvironmentID: env.ID, AllocationID: owner.ID}, ProviderID: owner.ID, State: "running", BootstrapComplete: true} - reconcileManagedState(t, w, db, tenant, env.ID, "released") - got, err = fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + reconcileManagedState(t, w, s, tenant, env.ID, "released") + got, err = deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "released" || len(p.resources) != 0 { t.Fatalf("late creation escaped cleanup: %+v %v", got, err) } } func TestManagedRuntimeExpiryRevokesWhenProviderUnavailable(t *testing.T) { - s, db := newManagedTestStoreDB(t) - tenant, _, env := managedSession(t, s, db) + s, _ := newManagedTestStore(t) + tenant, _, env := managedSession(t, s) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} - w, _ := managedWorker(t, s, db, key, p) + w, _ := managedWorker(t, s, key, p) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) } - if _, err := db.pool.Exec(t.Context(), "UPDATE runtime_allocations SET kept_at=clock_timestamp()-interval '61 minutes' WHERE id=$1", owner.ID); err != nil { + if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_allocations SET kept_at=clock_timestamp()-interval '61 minutes' WHERE id=$1", owner.ID); err != nil { t.Fatal(err) } p.unavailable = true - reconcileManagedState(t, w, db, tenant, env.ID, "cleanup_pending") - got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + reconcileManagedState(t, w, s, tenant, env.ID, "cleanup_pending") + got, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "cleanup_pending" { t.Fatalf("expiry lost on provider failure: %+v %v", got, err) } - if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("expired credential still authenticates") } if p.kills != 0 { @@ -241,11 +241,11 @@ func TestManagedRuntimeExpiryRevokesWhenProviderUnavailable(t *testing.T) { } func TestManagedRuntimeStoppedComputeDoesNotRequestCleanup(t *testing.T) { - s, db := newManagedTestStoreDB(t) - tenant, _, env := managedSession(t, s, db) + s, _ := newManagedTestStore(t) + tenant, _, env := managedSession(t, s) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} - w, _ := managedWorker(t, s, db, key, p) + w, _ := managedWorker(t, s, key, p) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) @@ -266,7 +266,7 @@ func TestManagedRuntimeStoppedComputeDoesNotRequestCleanup(t *testing.T) { if p.gets == before { t.Fatal("fixture allocation not inspected") } - got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + got, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "running" || p.kills != 0 || p.creates != 1 { t.Fatalf("compute interruption authorized replacement/cleanup: %+v %v", got, err) } diff --git a/services/core/tests/integration/runtime_node_lifecycle_fixture_test.go b/services/core/tests/integration/runtime_node_lifecycle_fixture_test.go index 0a1ed7040..d581b19f6 100644 --- a/services/core/tests/integration/runtime_node_lifecycle_fixture_test.go +++ b/services/core/tests/integration/runtime_node_lifecycle_fixture_test.go @@ -77,7 +77,6 @@ func (p *nodeIsolationProvider) RunCommand(ctx context.Context, r sandbox.Refere type nodeIsolationFixture struct { t *testing.T store *Store - db fixtureDB nodes *deployment.Service pool *pgxpool.Pool worker *execution.Worker @@ -97,8 +96,7 @@ func newNodeIsolationFixture(t *testing.T, mode string) *nodeIsolationFixture { 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) registry := runtimegateway.NewRegistry() cp := &fakeCheckpointProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, computes: map[string]sandbox.ComputeState{}, snapshots: map[string]sandbox.SnapshotIdentity{}, bootstraps: map[string]sandbox.Bootstrap{}, peers: map[string]*websocket.Conn{}, registry: registry} p := &nodeIsolationProvider{fakeCheckpointProvider: cp, blocked: map[string]bool{}, mode: mode, entered: make(chan struct{})} @@ -115,7 +113,7 @@ func newNodeIsolationFixture(t *testing.T, mode string) *nodeIsolationFixture { p.writes.Add(1) return completedInitialization(request, data) }} - handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(fixtureSessionStore(db)), Registry: registry}) + handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(sessionAdapter(s)), Registry: registry}) server := httptest.NewServer(http.HandlerFunc(handler.WS)) cp.endpoint = "ws" + strings.TrimPrefix(server.URL, "http") t.Cleanup(func() { @@ -126,11 +124,11 @@ func newNodeIsolationFixture(t *testing.T, mode string) *nodeIsolationFixture { cp.mu.Unlock() server.Close() }) - f := &nodeIsolationFixture{initializationCancel: cancelPreparation, t: t, store: s, db: db, nodes: fixtureDeployment(t, db), pool: pool, provider: p, key: uuid.NewString(), nodeA: uuid.NewString(), nodeB: uuid.NewString()} + f := &nodeIsolationFixture{initializationCancel: cancelPreparation, t: t, store: s, nodes: deploymentService(t, s), pool: pool, provider: p, key: uuid.NewString(), nodeA: uuid.NewString(), nodeB: uuid.NewString()} // Keep restored compute awake throughout the isolation assertions. // The suspension setup explicitly dates its activity two minutes in the past. policy := &execution.RuntimeSuspensionPolicy{IdleTimeout: time.Minute, Retention: time.Hour, MaxActive: 100, MaxRetained: 100} - f.worker = startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: f.key, BackendFingerprint: strings.Repeat("a", 64), Provider: p, ProviderKind: "microsandbox", LocalNodeID: f.nodeA, LocalCredentialSHA256: runtimedevice.HashCredential("local-credential"), LocalMaxActive: 100, LocalMaxRetained: 100, Suspension: policy}}) + f.worker = startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: f.key, BackendFingerprint: strings.Repeat("a", 64), Provider: p, ProviderKind: "microsandbox", LocalNodeID: f.nodeA, LocalCredentialSHA256: runtimedevice.HashCredential("local-credential"), LocalMaxActive: 100, LocalMaxRetained: 100, Suspension: policy}}) spec := SandboxDeploymentTestSpec("microsandbox") raw, _ := json.Marshal(spec) if _, err := pool.Exec(t.Context(), "UPDATE runtime_deployment SET specification=$1", raw); err != nil { @@ -140,7 +138,7 @@ func newNodeIsolationFixture(t *testing.T, mode string) *nodeIsolationFixture { t.Fatal(err) } t.Cleanup(f.stop) - f.epoch = fixtureOwnerEpoch(t, db) + f.epoch = fixtureOwnerEpoch(t, s) f.enroll(f.nodeB) f.online(f.nodeA) f.online(f.nodeB) @@ -211,7 +209,7 @@ func (f *nodeIsolationFixture) session(node string, initialize bool) (string, se if err != nil { f.t.Fatal(err) } - env, err := sessionReads(f.pool).GetSessionEnvironment(f.t.Context(), tenant, session.ID) + env, err := sessionAdapter(f.store).GetSessionEnvironment(f.t.Context(), tenant, session.ID) if err != nil { f.t.Fatal(err) } @@ -237,7 +235,7 @@ func (f *nodeIsolationFixture) phase(tenant, environment, phase string) deployme if err := f.worker.ReconcileManagedRuntimes(f.t.Context()); err != nil { f.t.Fatal(err) } - owner, err := fixtureReader(f.db).EnvironmentAllocation(f.t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) + owner, err := deploymentStore(f.store).EnvironmentAllocation(f.t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) if err != nil { f.t.Fatal(err) } diff --git a/services/core/tests/integration/runtime_node_lifecycle_test.go b/services/core/tests/integration/runtime_node_lifecycle_test.go index 39bef5925..545e87866 100644 --- a/services/core/tests/integration/runtime_node_lifecycle_test.go +++ b/services/core/tests/integration/runtime_node_lifecycle_test.go @@ -48,7 +48,7 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { case <-time.After(8 * time.Second): t.Fatalf("node A did not enter blocked %s operation", mode) } - online, err := fixtureReader(f.db).NodeOnline(t.Context(), f.nodeA) + online, err := deploymentStore(f.store).NodeOnline(t.Context(), f.nodeA) if err != nil || !online { t.Fatal("test node must stay online while its helper is blocked", err) } @@ -108,14 +108,14 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { // Only independent normal five-second loops drive these transitions. // No manual reconciliation or wake hint accelerates healthy nodes. waitNodeIsolation(t, 18*time.Second, func() (bool, string) { - wake, e1 := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: wakeTenant, EnvironmentID: wakeEnv.ID}) - deleted, e2 := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: deleteTenant, EnvironmentID: deleteEnv.ID}) - initialized, e3 := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + wake, e1 := deploymentStore(f.store).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: wakeTenant, EnvironmentID: wakeEnv.ID}) + deleted, e2 := deploymentStore(f.store).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: deleteTenant, EnvironmentID: deleteEnv.ID}) + initialized, e3 := deploymentStore(f.store).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) f.provider.mu.Lock() restores := f.provider.restores f.provider.mu.Unlock() - state := fmt.Sprintf("wake=%s/%s err=%v; deleted=%s/%s err=%v; initialized=%s/%s/%s err=%v; restores=%d writes=%d blocked_returns=%d", wake.State, wake.ComputePhase, e1, deleted.State, deleted.ComputePhase, e2, initialized.State, initializationState(t, f.pool, initialized.TenantID, initialized.EnvironmentID), initialized.ComputePhase, e3, restores, f.provider.writes.Load(), f.provider.returned.Load()) - return e1 == nil && e2 == nil && e3 == nil && wake.ComputePhase == "running" && deleted.State == "released" && initializationState(t, f.pool, initialized.TenantID, initialized.EnvironmentID) == "complete" && initialized.ComputePhase == "running", state + state := fmt.Sprintf("wake=%s/%s err=%v; deleted=%s/%s err=%v; initialized=%s/%s/%s err=%v; restores=%d writes=%d blocked_returns=%d", wake.State, wake.ComputePhase, e1, deleted.State, deleted.ComputePhase, e2, initialized.State, initializationState(t, f.store, initialized.TenantID, initialized.EnvironmentID), initialized.ComputePhase, e3, restores, f.provider.writes.Load(), f.provider.returned.Load()) + return e1 == nil && e2 == nil && e3 == nil && wake.ComputePhase == "running" && deleted.State == "released" && initializationState(t, f.store, initialized.TenantID, initialized.EnvironmentID) == "complete" && initialized.ComputePhase == "running", state }) if f.provider.returned.Load() != 0 || f.provider.writes.Load() != 1 { t.Fatalf("A returned early or initialization replayed: returned=%d writes=%d", f.provider.returned.Load(), f.provider.writes.Load()) @@ -140,7 +140,7 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { t.Fatal(err) } waitNodeIsolation(t, 7*time.Second, func() (bool, string) { - owner, err := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: ct, EnvironmentID: ce.ID}) + owner, err := deploymentStore(f.store).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: ct, EnvironmentID: ce.ID}) return err == nil && owner.State == "released", fmt.Sprintf("new node allocation=%s/%s err=%v", owner.State, owner.ComputePhase, err) }) if err := f.nodes.RemoveNode(t.Context(), nodeC); err != nil { diff --git a/services/core/tests/integration/runtime_observation_scan_test.go b/services/core/tests/integration/runtime_observation_scan_test.go index 17a5d317a..b535ccf30 100644 --- a/services/core/tests/integration/runtime_observation_scan_test.go +++ b/services/core/tests/integration/runtime_observation_scan_test.go @@ -8,14 +8,14 @@ import ( ) func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) var expected []string for range 5 { - _, session, _ := managedSession(t, s, db) + _, session, _ := managedSession(t, s) expected = append(expected, session.ID) } slices.Sort(expected) - if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: sessionTenant(t, db, expected[2]), SessionID: expected[2]}); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: sessionTenant(t, s, expected[2]), SessionID: expected[2]}); err != nil { t.Fatal(err) } expected = append(expected[:2], expected[3:]...) @@ -23,7 +23,7 @@ func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *test var got []string cursor := "" for { - page, err := fixtureReader(db).ObservationSessions(t.Context(), cursor, 2) + page, err := deploymentStore(s).ObservationSessions(t.Context(), cursor, 2) if err != nil { t.Fatal(err) } @@ -47,16 +47,16 @@ func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *test if !slices.Equal(got, expected) { t.Fatalf("observation scan = %v, want %v", got, expected) } - if _, err := fixtureReader(db).ObservationSessions(t.Context(), "", 0); err == nil { + if _, err := deploymentStore(s).ObservationSessions(t.Context(), "", 0); err == nil { t.Fatal("zero observation page size was accepted") } } -func sessionTenant(t *testing.T, db fixtureDB, sessionID string) string { +func sessionTenant(t *testing.T, s *Store, sessionID string) string { t.Helper() // The deployment-wide scan intentionally discovers tenant identity without // enumerating configured API keys. Use that same read to locate this fixture. - page, err := fixtureReader(db).ObservationSessions(t.Context(), "", 100) + page, err := deploymentStore(s).ObservationSessions(t.Context(), "", 100) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/runtime_pending_test.go b/services/core/tests/integration/runtime_pending_test.go index ad25d023f..65128f89c 100644 --- a/services/core/tests/integration/runtime_pending_test.go +++ b/services/core/tests/integration/runtime_pending_test.go @@ -14,20 +14,20 @@ import ( ) func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) { - s, db := newManagedTestStoreDB(t) - tenant, idle, idleEnvironment := managedSession(t, s, db) + s, _ := newManagedTestStore(t) + tenant, idle, idleEnvironment := managedSession(t, s) initial, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"model":"test"},"environment":{"type":"openai_hosted"}}`), InitialInputs: []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"hello"}`)}}}) if err != nil { t.Fatal(err) } - _, deleted, deletedEnvironment := managedSession(t, s, db) + _, deleted, deletedEnvironment := managedSession(t, s) if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: deleted.TenantID, SessionID: deleted.ID}); err != nil { t.Fatal(err) } key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} start := func() *execution.Worker { - w := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p}}) + w := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p}}) return w } stop := func(w *execution.Worker) { @@ -48,15 +48,15 @@ func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) t.Fatal(err) } } - idleOwner, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: idleEnvironment.ID}) + idleOwner, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: idleEnvironment.ID}) if err != nil || idleOwner.State != "running" { t.Fatal("idle creation was stranded", idleOwner, err) } - initialOwner, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: initial.Environment.ID}) + initialOwner, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: initial.Environment.ID}) if err != nil || initialOwner.State != "running" { t.Fatal("initial creation was stranded", initialOwner, err) } - if _, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: deleted.TenantID, EnvironmentID: deletedEnvironment.ID}); err == nil { + if _, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: deleted.TenantID, EnvironmentID: deletedEnvironment.ID}); err == nil { t.Fatal("deleted Session provisioned") } waiting, err := sessionAdapter(s).GetSession(t.Context(), tenant, initial.ID) @@ -81,7 +81,7 @@ func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) t.Fatal("restart repeated bootstrap", creates, p.creates) } for _, owner := range []deployment.Allocation{idleOwner, initialOwner} { - got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: owner.EnvironmentID}) + got, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: owner.EnvironmentID}) if err != nil || got.ID != owner.ID || got.DeviceID != owner.DeviceID { t.Fatal("restart replaced allocation identity", got, err) } diff --git a/services/core/tests/integration/runtime_scan_test.go b/services/core/tests/integration/runtime_scan_test.go index c92f1a21b..c5221dbc1 100644 --- a/services/core/tests/integration/runtime_scan_test.go +++ b/services/core/tests/integration/runtime_scan_test.go @@ -28,13 +28,13 @@ func (p *scanProvider) GetInfo(ctx context.Context, ref sandbox.Reference) (sand func TestManagedRuntimeScanWrapServicesNextPage(t *testing.T) { for _, count := range []int{0, 1, 31, 32, 33, 65} { t.Run(fmt.Sprint(count), func(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) p := &scanProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}} key := uuid.NewString() - w, _ := managedWorker(t, s, db, key, p) + w, _ := managedWorker(t, s, key, p) var ids []string for range count { - tenant, _, env := managedSession(t, s, db) + tenant, _, env := managedSession(t, s) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) @@ -67,11 +67,11 @@ func TestManagedRuntimeScanWrapServicesNextPage(t *testing.T) { } func TestManagedRuntimeScanEmptyAfterCleanupAndCanceledCall(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) p := &scanProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}} key := uuid.NewString() - w, _ := managedWorker(t, s, db, key, p) - tenant, session, env := managedSession(t, s, db) + w, _ := managedWorker(t, s, key, p) + tenant, session, env := managedSession(t, s) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) @@ -96,7 +96,7 @@ func TestManagedRuntimeScanEmptyAfterCleanupAndCanceledCall(t *testing.T) { if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) + got, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "released" || p.kills != 1 { t.Fatal("cleanup delayed at EOF", got, err, p.kills) } @@ -110,7 +110,7 @@ func TestManagedRuntimeScanEmptyAfterCleanupAndCanceledCall(t *testing.T) { } } // A new allocation remains discoverable after the store becomes empty. - tenant, _, env = managedSession(t, s, db) + tenant, _, env = managedSession(t, s) next, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) diff --git a/services/core/tests/integration/runtime_wake_hint_integration_test.go b/services/core/tests/integration/runtime_wake_hint_integration_test.go index 318ff7474..f31a2e4a7 100644 --- a/services/core/tests/integration/runtime_wake_hint_integration_test.go +++ b/services/core/tests/integration/runtime_wake_hint_integration_test.go @@ -73,7 +73,7 @@ func newWakeHintIntegration(t *testing.T) *wakeHintIntegration { fakeCheckpointProvider: f.provider, sentinel: sentinel.owner.ID, release: make(chan struct{}), scans: make(chan int, 16), } - worker := startWorker(t, t.Context(), f.db, &execution.Dispatcher{ + worker := startWorker(t, t.Context(), f.store, &execution.Dispatcher{ Registry: f.provider.registry, ManagedRuntimes: &execution.RuntimeProvider{ CoreURL: "http://core.invalid/api/v1", InstallationID: f.key, @@ -121,7 +121,7 @@ func (f *wakeHintIntegration) pending(t *testing.T, target wakeHintIntegrationTa t.Helper() var id string awaitDaemonRemoteCondition(t, t.Context(), 2*time.Second, "committed wake input", func() bool { - return f.fixture.db.pool.QueryRow(t.Context(), + return f.fixture.store.pool.QueryRow(t.Context(), "SELECT id::text FROM environment_input_reservations WHERE session_id=$1 AND idempotency_key=$2", target.session.ID, key).Scan(&id) == nil }) @@ -166,7 +166,7 @@ func TestRuntimeWakeHintCommittedSubmitResumesBeforeNormalTick(t *testing.T) { } f.release() awaitDaemonRemoteCondition(t, t.Context(), 2*time.Second, "hint restored retained compute", func() bool { - owner, err := fixtureReader(f.fixture.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: f.target.tenant, EnvironmentID: f.target.environment.ID}) + owner, err := deploymentStore(f.fixture.store).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: f.target.tenant, EnvironmentID: f.target.environment.ID}) return err == nil && owner.ComputePhase == "running" }) if elapsed := time.Since(f.started); elapsed >= 3*time.Second { @@ -179,7 +179,7 @@ func TestRuntimeWakeHintCommittedSubmitResumesBeforeNormalTick(t *testing.T) { t.Fatal("wake replayed creation/restoration or sent native input", restores, creates, f.provider.promptFrames.Load()) } var reservations, turns int - if err := f.fixture.db.pool.QueryRow(t.Context(), + if err := f.fixture.store.pool.QueryRow(t.Context(), "SELECT (SELECT count(*) FROM environment_input_reservations WHERE session_id=$1), (SELECT count(*) FROM turns WHERE session_id=$1)", f.target.session.ID).Scan(&reservations, &turns); err != nil || reservations != 1 || turns != 1 { t.Fatal("retry duplicated input or started a Turn before preparation", reservations, turns, err) diff --git a/services/core/tests/integration/runtime_worker_recovery_test.go b/services/core/tests/integration/runtime_worker_recovery_test.go index 4e41fe586..654f8a444 100644 --- a/services/core/tests/integration/runtime_worker_recovery_test.go +++ b/services/core/tests/integration/runtime_worker_recovery_test.go @@ -60,7 +60,7 @@ func TestWorkerWaitsForComputeAndSurvivesPromotionConflict(t *testing.T) { if err != nil { t.Fatal(err) } - worker := startWorker(t, t.Context(), h.db, h.d) + worker := startWorker(t, t.Context(), h.s, h.d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -73,7 +73,7 @@ func TestWorkerWaitsForComputeAndSurvivesPromotionConflict(t *testing.T) { } }) frames := workerFrames(t, h) - environment, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) + environment, err := sessionAdapter(h.s).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) if err != nil { t.Fatal(err) } @@ -144,7 +144,7 @@ func TestWorkerRestartPreservesQueuedTurnWhileComputeWakes(t *testing.T) { if _, err := pool.Exec(t.Context(), `INSERT INTO turns(id,session_id,status) VALUES($1,$2,'queued')`, turn, h.session.ID); err != nil { t.Fatal(err) } - worker, err := startWorkerErr(t.Context(), h.db, h.d) + worker, err := startWorkerErr(t.Context(), h.s, h.d) if err != nil { t.Fatal("queued wake blocked Core startup", err) } diff --git a/services/core/tests/integration/sandbox_deployment_resources_test.go b/services/core/tests/integration/sandbox_deployment_resources_test.go index 8c256bf6b..f4588ac29 100644 --- a/services/core/tests/integration/sandbox_deployment_resources_test.go +++ b/services/core/tests/integration/sandbox_deployment_resources_test.go @@ -16,7 +16,7 @@ func TestSandboxDeploymentMutationViewsIncludeActualResources(t *testing.T) { if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) w := executionWriter(t, s) changes := deploymentExecution(t, w) installation := uuid.NewString() diff --git a/services/core/tests/integration/sandbox_deployment_switch_test.go b/services/core/tests/integration/sandbox_deployment_switch_test.go index 9931d2296..49ee8ba6a 100644 --- a/services/core/tests/integration/sandbox_deployment_switch_test.go +++ b/services/core/tests/integration/sandbox_deployment_switch_test.go @@ -37,7 +37,7 @@ func TestSandboxResetClearsCustomE2BEndpoint(t *testing.T) { if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) w := executionWriter(t, s) changes := deploymentExecution(t, w) installation := uuid.NewString() @@ -78,7 +78,7 @@ func TestSandboxDirectDeploymentOwnershipAndCleanSwitch(t *testing.T) { if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) w := executionWriter(t, s) changes := deploymentExecution(t, w) id := uuid.NewString() @@ -164,7 +164,7 @@ func TestSandboxDirectDeploymentOwnershipAndCleanSwitch(t *testing.T) { func TestSandboxSwitchRetiresNodesAndEnrollment(t *testing.T) { _, pool := newManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{5}, 32)) - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) w := executionWriter(t, s) changes := deploymentExecution(t, w) nodes := deploymentService(t, s) @@ -239,7 +239,7 @@ func TestSandboxSwitchRetiresNodesAndEnrollment(t *testing.T) { func TestSandboxResetSerializesFreshDirectSessions(t *testing.T) { _, pool := newManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{6}, 32)) - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) w := executionWriter(t, s) changes := deploymentExecution(t, w) id := uuid.NewString() @@ -283,7 +283,7 @@ func TestSandboxResetSerializesFreshDirectSessions(t *testing.T) { func TestSandboxSwitchPreservesReleasedAllocationAndItemHistory(t *testing.T) { _, pool := newManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{8}, 32)) - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) w := executionWriter(t, s) changes := deploymentExecution(t, w) installation := uuid.NewString() @@ -368,7 +368,7 @@ func TestSandboxSwitchPreservesReleasedAllocationAndItemHistory(t *testing.T) { func TestUnspecifiedNodeDeploymentRejectedWithoutMutation(t *testing.T) { _, pool := newManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{7}, 32)) - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) w := executionWriter(t, s) changes := deploymentExecution(t, w) nodes := deploymentService(t, s) diff --git a/services/core/tests/integration/sandbox_deployment_switch_worker_test.go b/services/core/tests/integration/sandbox_deployment_switch_worker_test.go index c8c21589d..e43dd7ae3 100644 --- a/services/core/tests/integration/sandbox_deployment_switch_worker_test.go +++ b/services/core/tests/integration/sandbox_deployment_switch_worker_test.go @@ -24,9 +24,8 @@ import ( func TestSandboxWorkerSwitchesAndRecoversFailedActivation(t *testing.T) { _, pool := newManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{7}, 32)) - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} - s.SetPlacement(fixtureRules(t, db)) - deployments := fixtureDeployment(t, db) + s := NewWithCredentialCipher(pool, cipher) + deployments := deploymentService(t, s) id := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} var fail atomic.Bool @@ -44,7 +43,7 @@ func TestSandboxWorkerSwitchesAndRecoversFailedActivation(t *testing.T) { } return execution.PreparedRuntimeDeployment{Config: &execution.RuntimeProvider{InstallationID: setup.InstallationID, ProviderKind: setup.Provider, Mode: setup.Mode, AdmissionPaused: setup.AdmissionPaused, CoreURL: "https://core.example/api/v1", BackendFingerprint: setup.BackendFingerprint, Provider: p}}, nil }) - w := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) + w := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- w.Run(ctx) }() @@ -146,7 +145,7 @@ func TestSandboxWorkerSwitchesAndRecoversFailedActivation(t *testing.T) { if _, err := w.InitializeSandboxDeployment(t.Context(), cloud); err != nil { t.Fatal("setup after unconfigured publication", err) } - tenant, session, environment := managedSession(t, s, db) + tenant, session, environment := managedSession(t, s) allocation, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, id) if err != nil || allocation.NodeID != "" || allocation.State != "running" { t.Fatal("direct provider not available after resume", allocation, err) @@ -158,7 +157,7 @@ func TestSandboxWorkerSwitchesAndRecoversFailedActivation(t *testing.T) { if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } - reconcileManagedState(t, w, db, tenant, environment.ID, "released") + reconcileManagedState(t, w, s, tenant, environment.ID, "released") empty = reset(4) next.ExpectedGeneration = empty.Generation if _, err := w.InitializeSandboxDeployment(t.Context(), next); err != nil { diff --git a/services/core/tests/integration/sandbox_deployment_view_test.go b/services/core/tests/integration/sandbox_deployment_view_test.go index 000a058df..0e696c890 100644 --- a/services/core/tests/integration/sandbox_deployment_view_test.go +++ b/services/core/tests/integration/sandbox_deployment_view_test.go @@ -20,7 +20,7 @@ func TestSandboxDeploymentViewRecordsTemplateBuildAndSuspension(t *testing.T) { if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) w := executionWriter(t, s) changes := deploymentExecution(t, w) id := uuid.NewString() diff --git a/services/core/tests/integration/sandbox_deployment_worker_test.go b/services/core/tests/integration/sandbox_deployment_worker_test.go index 4b28da8d1..95538e64a 100644 --- a/services/core/tests/integration/sandbox_deployment_worker_test.go +++ b/services/core/tests/integration/sandbox_deployment_worker_test.go @@ -18,8 +18,8 @@ import ( ) func TestSandboxDeploymentWorkerActivatesWithoutRestart(t *testing.T) { - s, db := newManagedTestStoreDB(t) - deployments := fixtureDeployment(t, db) + s, _ := newManagedTestStore(t) + deployments := deploymentService(t, s) id := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} configuration := execution.NewDeferredRuntimeProvider(id, func(ctx context.Context) (*execution.RuntimeProvider, error) { @@ -34,7 +34,7 @@ func TestSandboxDeploymentWorkerActivatesWithoutRestart(t *testing.T) { }) start := func() (*execution.Worker, func()) { t.Helper() - w := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) + w := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) var once sync.Once stop := func() { once.Do(func() { ctx, cancel := context.WithCancel(context.Background()); cancel(); _ = w.Run(ctx) }) @@ -63,7 +63,7 @@ func TestSandboxDeploymentWorkerActivatesWithoutRestart(t *testing.T) { } connect := func() { t.Helper() - epoch := fixtureOwnerEpoch(t, db) + epoch := fixtureOwnerEpoch(t, s) connection := uuid.NewString() if err := deployments.ConnectNode(t.Context(), nodeID, connection, epoch); err != nil { t.Fatal(err) @@ -73,7 +73,7 @@ func TestSandboxDeploymentWorkerActivatesWithoutRestart(t *testing.T) { } } connect() - tenant, _, environment := managedSession(t, s, db) + tenant, _, environment := managedSession(t, s) allocation, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, id) if err != nil || allocation.NodeID != nodeID || p.creates != 1 { t.Fatal("activation failed", allocation, err) diff --git a/services/core/tests/integration/sandbox_node_auth_order_http_test.go b/services/core/tests/integration/sandbox_node_auth_order_http_test.go index 5e92a1965..67f9d9615 100644 --- a/services/core/tests/integration/sandbox_node_auth_order_http_test.go +++ b/services/core/tests/integration/sandbox_node_auth_order_http_test.go @@ -17,19 +17,19 @@ import ( // authenticate first: a missing or invalid credential gets 401, and only a // recognized credential learns that the deployment is unavailable. func TestSandboxNodeRoutesAuthenticateBeforeDeploymentState(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential(uuid.NewString())}) if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, db), withCoreKeys(admin)) + handler, err := publicHandler(t, s, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, s), withCoreKeys(admin)) if err != nil { t.Fatal(err) } // Recognized, unconsumed enrollment tokens; no deployment has been initialized. enrollment := func(token, installation string) { t.Helper() - if _, err := db.pool.Exec(t.Context(), "INSERT INTO runtime_node_enrollments(token_sha256,installation_id,expires_at) VALUES(encode(sha256($1::bytea),'hex'),$2,clock_timestamp()+interval '10 minutes')", token, installation); err != nil { + if _, err := s.pool.Exec(t.Context(), "INSERT INTO runtime_node_enrollments(token_sha256,installation_id,expires_at) VALUES(encode(sha256($1::bytea),'hex'),$2,clock_timestamp()+interval '10 minutes')", token, installation); err != nil { t.Fatal(err) } } @@ -75,7 +75,7 @@ func TestSandboxNodeRoutesAuthenticateBeforeDeploymentState(t *testing.T) { // Once Web claims an installation, still before initialization, another // installation's token gets the same 401 it gets after initialization. - if _, err := db.pool.Exec(t.Context(), "UPDATE runtime_deployment SET installation_id=$1, web_managed=true WHERE singleton=true", claimed); err != nil { + if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_deployment SET installation_id=$1, web_managed=true WHERE singleton=true", claimed); err != nil { t.Fatal(err) } run([]check{ diff --git a/services/core/tests/integration/sandbox_specification_lifecycle_test.go b/services/core/tests/integration/sandbox_specification_lifecycle_test.go index baf925f1c..19e9e2513 100644 --- a/services/core/tests/integration/sandbox_specification_lifecycle_test.go +++ b/services/core/tests/integration/sandbox_specification_lifecycle_test.go @@ -24,7 +24,7 @@ func webSpecificationFixture(t *testing.T, provider string) (*Store, *Store, dep if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) w := executionWriter(t, s) id := uuid.NewString() changes := deploymentExecution(t, w) diff --git a/services/core/tests/integration/saved_web_search_public_test.go b/services/core/tests/integration/saved_web_search_public_test.go index 96992c293..dd31102ac 100644 --- a/services/core/tests/integration/saved_web_search_public_test.go +++ b/services/core/tests/integration/saved_web_search_public_test.go @@ -19,13 +19,13 @@ import ( // and keeps tenant isolation (W8). func TestSavedWebSearchPostgres(t *testing.T) { // An isolated database keeps the no-write digest independent of other tests. - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) owner, foreign, ownerTenant := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "search-owner", TokenSHA256: runtimedevice.HashCredential(owner), TenantID: ownerTenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "search-foreign", 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) } @@ -115,7 +115,7 @@ func TestSavedWebSearchPostgres(t *testing.T) { // A disabled record saved before this batch reads unchanged and is admitted with // the same frozen Session tool (W7). - _, agentService := fixtureAgents(t, db) + _, agentService := fixtureAgents(t, s) legacy, err := agentService.Create(t.Context(), agents.CreateCommand{TenantID: ownerTenant, Metadata: map[string]string{}, Configuration: json.RawMessage( `{"model":"search-model","name":null,"instructions":null,"multi_agent":{"enabled":false,"max_concurrent_subagents":null},"reasoning":{},"service_tier":"auto","text":{"format":{"type":"text"},"verbosity":"medium"},"tools":[{"type":"web_search","mode":"disabled","context_size":"medium","allowed_domains":[],"location":null}]}`)}) if err != nil { @@ -172,7 +172,7 @@ func TestSavedWebSearchPostgres(t *testing.T) { enabledTools := map[string]string{"mode-live": agentIDs["mode-live"], "mode-cached": agentIDs["mode-cached"], "type-only": agentIDs["type-only"], "updated-to-live": agentIDs["mode-disabled"]} // W4: every creation mode rejects enabled saved search without writes. - before := databaseDigest(t, db.pool) + before := databaseDigest(t, s.pool) const rejection = `{"error":{"message":"Only disabled web_search is qualified for execution.","type":"invalid_request_error","code":"unsupported_or_invalid_configuration","param":null}}` + "\n" for name, id := range enabledTools { for _, suffix := range []string{ @@ -224,7 +224,7 @@ func TestSavedWebSearchPostgres(t *testing.T) { if status, raw := client.do(foreign, http.MethodGet, "/v1/agents?limit=100", "", nil); status != http.StatusOK || strings.Contains(raw, "search-model") { t.Errorf("foreign list: %d %s", status, raw) } - if after := databaseDigest(t, db.pool); !mapsEqual(before, after) { + if after := databaseDigest(t, s.pool); !mapsEqual(before, after) { t.Fatal("rejected Session creation changed persisted state") } diff --git a/services/core/tests/integration/self_hosted_cancel_public_test.go b/services/core/tests/integration/self_hosted_cancel_public_test.go index ae60019a1..31f6f8989 100644 --- a/services/core/tests/integration/self_hosted_cancel_public_test.go +++ b/services/core/tests/integration/self_hosted_cancel_public_test.go @@ -23,7 +23,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, db := newModelTestStoreDB(t) + s, _ := NewModelTestStore(t) tenant, foreignTenant := uuid.NewString(), uuid.NewString() token, foreign := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ @@ -33,9 +33,9 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { var owner execution.Owner serve := func() (*httptest.Server, func(bool)) { t.Helper() - owner = executionOwner(t, db) - worker, stop := publicOwnedWorker(t, s, db, owner) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, worker), executorURL("https://offline-executor.example")) + owner = executionOwner(t, s) + worker, stop := publicOwnedWorker(t, s, owner) + handler, err := publicHandler(t, s, auth, "codex", workerExecution(t, worker), executorURL("https://offline-executor.example")) if err != nil { t.Fatal(err) } @@ -79,7 +79,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { settings["accepted"] = accepted receipts := func(key, target string) []sessions.InputReceipt { t.Helper() - rows, err := db.pool.Query(t.Context(), `SELECT sequence, COALESCE(turn_id::text,'') FROM turn_inputs + rows, err := s.pool.Query(t.Context(), `SELECT sequence, COALESCE(turn_id::text,'') FROM turn_inputs WHERE session_id=$1 AND idempotency_key=$2 ORDER BY batch_position`, created.ID, key) if err != nil { t.Fatal(err) @@ -104,7 +104,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { snapshot := func(sessionID string) string { t.Helper() var value string - err := db.pool.QueryRow(t.Context(), `SELECT jsonb_build_object( + err := s.pool.QueryRow(t.Context(), `SELECT jsonb_build_object( 'session', (SELECT to_jsonb(s) FROM sessions s WHERE id=$1), 'reservations', (SELECT jsonb_agg(to_jsonb(r) ORDER BY r.id) FROM environment_input_reservations r WHERE session_id=$1), 'turns', (SELECT jsonb_agg(to_jsonb(t) ORDER BY t.id) FROM turns t WHERE session_id=$1), @@ -148,7 +148,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { if snapshot(created.ID) != before { t.Fatal("idle cancellation replay or rejected input changed active work") } - itemsBefore, err := sessionReads(db.pool).ListItems(t.Context(), tenant, created.ID, "", 100, true) + itemsBefore, err := sessionAdapter(s).ListItems(t.Context(), tenant, created.ID, "", 100, true) if err != nil || len(itemsBefore.Items) != 2 { t.Fatal("controlled partial output was not recorded", err) } @@ -162,7 +162,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { if err != nil || turn.Status != sessions.TurnInProgress || turn.CancelRequestedAt.IsZero() || !turn.CompletedAt.IsZero() { t.Fatal("202 must admit cancellation without fabricating native completion", err) } - itemsAfter, err := sessionReads(db.pool).ListItems(t.Context(), tenant, created.ID, "", 100, true) + itemsAfter, err := sessionAdapter(s).ListItems(t.Context(), tenant, created.ID, "", 100, true) if err != nil || !reflect.DeepEqual(itemsBefore, itemsAfter) { t.Fatal("cancellation admission changed partial history", err) } @@ -173,12 +173,12 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { transition(first, sessions.TurnInProgress, sessions.TurnCancelled) for _, reopen := range []bool{false, true} { if reopen { - awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, db.pool) + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, s.pool) stop(false) awaitRelease() server.Close() - db.pool.Close() - s, db = newModelTestStoreDB(t) + s.pool.Close() + s, _ = NewModelTestStore(t) server, stop = serve() settings["base"] = server.URL } diff --git a/services/core/tests/integration/self_hosted_initial_public_test.go b/services/core/tests/integration/self_hosted_initial_public_test.go index 0693f210a..be5360247 100644 --- a/services/core/tests/integration/self_hosted_initial_public_test.go +++ b/services/core/tests/integration/self_hosted_initial_public_test.go @@ -30,7 +30,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, db := newModelTestStoreDB(t) + s, _ := NewModelTestStore(t) tenant, foreignTenant := uuid.NewString(), uuid.NewString() token, peer, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ @@ -39,7 +39,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { {OrganizationID: "test-org", ProjectID: foreignTenant, SubjectKind: "service_account", SubjectID: "initial-creator", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: foreignTenant}, }) const origin = "https://offline-executor.example" - serve := func(s *Store, db fixtureDB, worker *execution.Worker) *httptest.Server { + serve := func(s *Store, worker *execution.Worker) *httptest.Server { t.Helper() enabled := []func(*api.Dependencies){acceptUnavailable(t)} if worker != nil { @@ -56,7 +56,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { } }) } - handler, err := publicHandler(t, s, db, auth, "codex", enabled...) + handler, err := publicHandler(t, s, auth, "codex", enabled...) if err != nil { t.Fatal(err) } @@ -64,8 +64,8 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { t.Cleanup(server.Close) return server } - worker, stop := publicInitialWorker(t, s, db) - server := serve(s, db, worker) + worker, stop := publicInitialWorker(t, s) + server := serve(s, worker) settings := map[string]any{"base": server.URL, "token": token, "peer_token": peer, "foreign_token": foreign, "remote_url": origin} run := func(phase string) json.RawMessage { t.Helper() @@ -129,7 +129,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { if !reflect.DeepEqual(texts, item.Texts) { t.Fatal("public initial text order changed") } - environment, err := sessionReads(pool).GetSessionEnvironment(t.Context(), tenant, item.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, item.ID) if err != nil || environment.ID != item.EnvironmentID || environment.Status != "pending" { t.Fatal("public initial Environment identity changed", err) } @@ -144,23 +144,23 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { } return result } - before := reservations(s, db.pool) + before := reservations(s, s.pool) for _, reservation := range before { if reservation.State != sessions.EnvironmentInputPending || reservation.Deadline.Sub(reservation.CreatedAt) != 5*time.Minute { t.Fatal("public initial creation did not retain its database deadline") } } - awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, db.pool) + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, s.pool) stop(false) awaitRelease() server.Close() - db.pool.Close() - reopened, reopenedDB := newModelTestStoreDB(t) - worker, stop = publicInitialWorker(t, reopened, reopenedDB) - server = serve(reopened, reopenedDB, worker) + s.pool.Close() + reopened, _ := NewModelTestStore(t) + worker, stop = publicInitialWorker(t, reopened) + server = serve(reopened, worker) settings["base"], settings["accepted"] = server.URL, accepted run("reopen") - if !reflect.DeepEqual(before, reservations(reopened, reopenedDB.pool)) { + if !reflect.DeepEqual(before, reservations(reopened, reopened.pool)) { t.Fatal("reopened public retry changed reservation identity or deadline") } failureID := created.Cases[0].ID @@ -170,7 +170,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { return } // Advance one known deadline; the running Worker still owns settlement and events. - tag, err := reopenedDB.pool.Exec(r.Context(), "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1 AND is_initial AND state='pending'", failureID) + tag, err := reopened.pool.Exec(r.Context(), "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1 AND is_initial AND state='pending'", failureID) if err != nil || tag.RowsAffected() != 1 { t.Error("controlled initial deadline update failed", err) w.WriteHeader(http.StatusInternalServerError) @@ -181,7 +181,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { defer control.Close() settings["expiry_control"] = control.URL run("expire") - after := reservations(reopened, reopenedDB.pool) + after := reservations(reopened, reopened.pool) for id, reservation := range after { if id == failureID { if reservation.ID != before[id].ID || reservation.State != sessions.EnvironmentInputExpired || reservation.SettledAt == nil { @@ -192,7 +192,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { } } var pid uint32 - err := reopenedDB.pool.QueryRow(t.Context(), `SELECT pid FROM pg_locks WHERE locktype='advisory' + err := reopened.pool.QueryRow(t.Context(), `SELECT pid FROM pg_locks WHERE locktype='advisory' AND database=(SELECT oid FROM pg_database WHERE datname=current_database()) AND classid=(706172736172::bigint >> 32)::oid AND objid=(706172736172::bigint & 4294967295)::oid AND objsubid=1 AND granted`).Scan(&pid) @@ -200,27 +200,27 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { t.Fatal(err) } var killed bool - if err := reopenedDB.pool.QueryRow(t.Context(), "SELECT pg_terminate_backend($1, 1000)", pid).Scan(&killed); err != nil || !killed { + if err := reopened.pool.QueryRow(t.Context(), "SELECT pg_terminate_backend($1, 1000)", pid).Scan(&killed); err != nil || !killed { t.Fatal("could not end the fixture Worker's execution lease", err) } stop(true) - settings["disabled_base"] = serve(reopened, reopenedDB, nil).URL + settings["disabled_base"] = serve(reopened, nil).URL run("unavailable") - if !reflect.DeepEqual(after, reservations(reopened, reopenedDB.pool)) { + if !reflect.DeepEqual(after, reservations(reopened, reopened.pool)) { t.Fatal("unavailable execution or recorded retry changed initial work") } } -func publicInitialWorker(t *testing.T, s *Store, db fixtureDB) (*execution.Worker, func(bool)) { +func publicInitialWorker(t *testing.T, s *Store) (*execution.Worker, func(bool)) { t.Helper() - return publicOwnedWorker(t, s, db, executionOwner(t, db)) + return publicOwnedWorker(t, s, executionOwner(t, s)) } // publicOwnedWorker is publicInitialWorker on owner, for tests that also write // as the Worker's execution owner. -func publicOwnedWorker(t *testing.T, s *Store, db fixtureDB, owner execution.Owner) (*execution.Worker, func(bool)) { +func publicOwnedWorker(t *testing.T, s *Store, owner execution.Owner) (*execution.Worker, func(bool)) { t.Helper() - worker := startOwnedWorker(t, t.Context(), db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry()}, owner) + worker := startOwnedWorker(t, t.Context(), s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry()}, owner) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() diff --git a/services/core/tests/integration/session_agent_filter_public_test.go b/services/core/tests/integration/session_agent_filter_public_test.go index d4fb2142c..e6239c311 100644 --- a/services/core/tests/integration/session_agent_filter_public_test.go +++ b/services/core/tests/integration/session_agent_filter_public_test.go @@ -17,20 +17,20 @@ func TestSessionAgentFilterOfficialClient(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) } diff --git a/services/core/tests/integration/session_artifacts_public_test.go b/services/core/tests/integration/session_artifacts_public_test.go index e8e3f61f1..4b72e8863 100644 --- a/services/core/tests/integration/session_artifacts_public_test.go +++ b/services/core/tests/integration/session_artifacts_public_test.go @@ -68,7 +68,7 @@ func completeArtifactTurn(t *testing.T, s *Store, artifacts *sessions.Service, t } // artifactHTTPServer serves Artifact routes for an owner and a foreign tenant. -func artifactHTTPServer(t *testing.T, s *Store, db fixtureDB) (server *httptest.Server, owner, ownerTenant, foreign, foreignTenant string) { +func artifactHTTPServer(t *testing.T, s *Store) (server *httptest.Server, owner, ownerTenant, foreign, foreignTenant string) { t.Helper() owner, foreign = uuid.NewString(), uuid.NewString() ownerTenant, foreignTenant = uuid.NewString(), uuid.NewString() @@ -76,7 +76,7 @@ func artifactHTTPServer(t *testing.T, s *Store, db fixtureDB) (server *httptest. {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "artifact-owner", TokenSHA256: runtimedevice.HashCredential(owner), TenantID: ownerTenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "artifact-foreign", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: foreignTenant}, }) - h, err := publicHandler(t, s, db, auth, "codex") + h, err := publicHandler(t, s, auth, "codex") if err != nil { t.Fatal(err) } @@ -89,12 +89,12 @@ func artifactHTTPServer(t *testing.T, s *Store, db fixtureDB) (server *httptest. // environment_id filter matches nothing like another Environment's ID (HE-56), // without weakening tenant or Session scoping. func TestSessionArtifactListEnvelopeAndEnvironmentFilterPostgres(t *testing.T) { - s, db := newTestStoreDB(t) - _, sessionService, err := fixtureSessions(db) + s, _ := testStore(t) + sessionService, err := newSessionService(s) if err != nil { t.Fatal(err) } - server, owner, ownerTenant, foreign, foreignTenant := artifactHTTPServer(t, s, db) + server, owner, ownerTenant, foreign, foreignTenant := artifactHTTPServer(t, s) client := pathIDClient{t: t, server: server} session, environment := hostedArtifactSession(t, s, ownerTenant, "artifact-list") @@ -213,12 +213,13 @@ func TestSessionArtifactsOfficialClientPostgres(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, db := newTestStoreDB(t) - sessionStore, sessionService, err := fixtureSessions(db) + s, _ := testStore(t) + sessionStore := sessionAdapter(s) + sessionService, err := newSessionService(s) if err != nil { t.Fatal(err) } - server, owner, ownerTenant, foreign, _ := artifactHTTPServer(t, s, db) + server, owner, ownerTenant, foreign, _ := artifactHTTPServer(t, s) session, environment := hostedArtifactSession(t, s, ownerTenant, "artifact-sdk") outputs := map[string]string{"a.txt": "alpha", "sub/b.txt": "bravo", "empty.txt": ""} first := completeArtifactTurn(t, s, sessionService, ownerTenant, session, environment, "artifact-sdk-1", outputs) diff --git a/services/core/tests/integration/session_deletion_execution_test.go b/services/core/tests/integration/session_deletion_execution_test.go index 8a6124e92..fb1872d74 100644 --- a/services/core/tests/integration/session_deletion_execution_test.go +++ b/services/core/tests/integration/session_deletion_execution_test.go @@ -15,7 +15,7 @@ func TestDeletedSessionWaitingTurnSettlesWithoutStoppingWorker(t *testing.T) { h := newFunctionHarness(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() { @@ -77,7 +77,7 @@ func TestDeletedSessionRestartStillReconcilesHiddenClaim(t *testing.T) { if err := h.s.commitLegacyDeletion(ctx, h.tenant, h.session.ID); err != nil { t.Fatal(err) } - worker := startWorker(t, ctx, h.db, h.d) + worker := startWorker(t, ctx, h.s, h.d) stopped, cancel := context.WithCancel(ctx) cancel() if err := worker.Run(stopped); !errors.Is(err, context.Canceled) { @@ -98,7 +98,7 @@ func TestWaitingSessionCancelsThenDeletesThroughWorker(t *testing.T) { h := newFunctionHarness(t) 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() { diff --git a/services/core/tests/integration/session_deletion_lifecycle_public_test.go b/services/core/tests/integration/session_deletion_lifecycle_public_test.go index 930530ec1..57216d5a3 100644 --- a/services/core/tests/integration/session_deletion_lifecycle_public_test.go +++ b/services/core/tests/integration/session_deletion_lifecycle_public_test.go @@ -23,22 +23,22 @@ const deletionAgent = `"agent":{"id":"agent_deletion","model":"fixture","tools": // missing and malformed identifiers keep one not-found response. func TestSessionDeletionLifecyclePostgres(t *testing.T) { // An isolated database keeps the no-write digest independent of other tests. - s, db := newManagedTestStoreDB(t) - audit := auditpg.New(pgunit.NewPool(db.pool)) + s, _ := newManagedTestStore(t) + audit := auditpg.New(pgunit.NewPool(s.pool)) ctx := t.Context() tenant, owner, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "deletion-owner", TokenSHA256: runtimedevice.HashCredential(owner), TenantID: tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "deletion-foreign", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) + h, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() client := pathIDClient{t: t, server: server} - leased := executionOwner(t, db) + leased := executionOwner(t, s) create := func(environment string, initial bool) sessions.Session { t.Helper() @@ -134,7 +134,7 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { "later_input_cancelled": "", } expired := create(selfHosted, true) - if _, err := db.pool.Exec(ctx, "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1", expired.ID); err != nil { + if _, err := s.pool.Exec(ctx, "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1", expired.ID); err != nil { t.Fatal(err) } if count, err := leased.Sessions.ExpireEnvironmentInputs(ctx); err != nil || count != 1 { @@ -181,13 +181,13 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { if readStatus != http.StatusOK || decode(before)["status"] != projected[name] { t.Fatal(readStatus, before) } - digest := databaseDigest(t, db.pool) + digest := databaseDigest(t, s.pool) notFound(foreign, http.MethodDelete, sessionPath(id)) status, raw := client.do(owner, http.MethodDelete, sessionPath(id), "", nil) if status != http.StatusConflict || !reflect.DeepEqual(decode(raw), conflict) { t.Fatalf("busy Session deletion: %d %s", status, raw) } - if after := databaseDigest(t, db.pool); !reflect.DeepEqual(after, digest) { + if after := databaseDigest(t, s.pool); !reflect.DeepEqual(after, digest) { t.Fatal("rejected deletion changed the database") } if readStatus, after := client.do(owner, http.MethodGet, sessionPath(id), "", nil); readStatus != http.StatusOK || after != before { @@ -202,7 +202,7 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { if status != http.StatusOK || !reflect.DeepEqual(decode(first), map[string]any{"id": id, "object": "agent.session.deleted", "deleted": true}) { t.Fatalf("settled Session deletion: %d %s", status, first) } - digest := databaseDigest(t, db.pool) + digest := databaseDigest(t, s.pool) filter := writeaudit.Filter{ResourceType: "session", ResourceID: id, Limit: 100} beforeAudit, err := audit.ListWriteOperations(ctx, tenant, filter) if err != nil { @@ -224,7 +224,7 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { } // Only the new operation records may differ. Ownership, Session state, // execution data and every public response remain unchanged. - after := databaseDigest(t, db.pool) + after := databaseDigest(t, s.pool) delete(after, "write_audit_operations") delete(digest, "write_audit_operations") if !reflect.DeepEqual(after, digest) { diff --git a/services/core/tests/integration/session_deletion_public_test.go b/services/core/tests/integration/session_deletion_public_test.go index 6b60be3bd..65fae5e91 100644 --- a/services/core/tests/integration/session_deletion_public_test.go +++ b/services/core/tests/integration/session_deletion_public_test.go @@ -17,19 +17,19 @@ func TestSessionDeletionOfficialClient(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() - h, err = publicHandler(t, New(db.pool), db, auth, "codex", storeExecution(t, New(db.pool))) + h, err = publicHandler(t, New(s.pool), auth, "codex", storeExecution(t, New(s.pool))) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/session_devices_fixture_test.go b/services/core/tests/integration/session_devices_fixture_test.go index f05c43240..9dc9380e5 100644 --- a/services/core/tests/integration/session_devices_fixture_test.go +++ b/services/core/tests/integration/session_devices_fixture_test.go @@ -14,13 +14,13 @@ import ( // bindSessionDevice binds the tenant's device to the Session through the // Session execution operations on an execution lease of its own, and returns // once the server released that lease, so a Worker can take it next. -func bindSessionDevice(t *testing.T, db fixtureDB, tenant, session, device string) error { +func bindSessionDevice(t *testing.T, s *Store, tenant, session, device string) error { t.Helper() - lease, err := pgunit.AcquireLease(t.Context(), db.pool) + lease, err := pgunit.AcquireLease(t.Context(), s.pool) if err != nil { return err } - released := pgtest.ObserveExecutionLeaseRelease(t, db.pool) + released := pgtest.ObserveExecutionLeaseRelease(t, s.pool) operations, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) if err == nil { err = operations.BindSessionDevice(t.Context(), tenant, session, device) diff --git a/services/core/tests/integration/session_events_test.go b/services/core/tests/integration/session_events_test.go index fc96e11b8..b68308acb 100644 --- a/services/core/tests/integration/session_events_test.go +++ b/services/core/tests/integration/session_events_test.go @@ -50,7 +50,7 @@ func TestSessionEventsAreVisibleOnlyAfterCommit(t *testing.T) { func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { ctx := context.Background() s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "stream"}) if err != nil { @@ -142,7 +142,7 @@ func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { t.Fatal("foreign event access", err) } before, _ = sessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) - if _, err = sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true); err != nil { + if _, err = sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 100, true); err != nil { t.Fatal(err) } after, _ = sessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) @@ -195,7 +195,7 @@ func TestSessionEventsRetentionAndQueuedCancellation(t *testing.T) { if _, err = pool.Exec(ctx, "UPDATE session_events SET payload=jsonb_build_object('padding',repeat('x',524288)) WHERE session_id=$1", session.ID); err != nil { t.Fatal(err) } - if _, err = sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 1, true); err != nil { + if _, err = sessionAdapter(s).ListItems(ctx, tenant, session.ID, "", 1, true); err != nil { t.Fatal(err) } var bytes int64 diff --git a/services/core/tests/integration/session_execution_configuration_test.go b/services/core/tests/integration/session_execution_configuration_test.go index f32f7e7bd..3f6f5f45a 100644 --- a/services/core/tests/integration/session_execution_configuration_test.go +++ b/services/core/tests/integration/session_execution_configuration_test.go @@ -37,7 +37,7 @@ func TestSessionExecutionConfigurationFrozenAcrossCreationPathsAndRetry(t *testi if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) for _, stream := range []bool{false, true} { for _, source := range []string{"session", "agent", "deployment"} { t.Run(source+map[bool]string{false: "/ordinary", true: "/stream"}[stream], func(t *testing.T) { @@ -165,7 +165,7 @@ func TestSessionExecutionConfigurationRollbackAndValidation(t *testing.T) { if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) tenant := uuid.NewString() for _, kind := range []string{"model", "harness", "source", "provider", "provider_mismatch", "post_projection_failure"} { input := executionProjectionInput("session") diff --git a/services/core/tests/integration/session_initial_public_test.go b/services/core/tests/integration/session_initial_public_test.go index f6e5a3f0d..3b8162891 100644 --- a/services/core/tests/integration/session_initial_public_test.go +++ b/services/core/tests/integration/session_initial_public_test.go @@ -17,11 +17,11 @@ func TestInitialSessionInputOfficialClient(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()}}) // Exercise real worker admission with dispatch paused for deterministic reads. - worker := startWorker(t, t.Context(), db, &execution.Dispatcher{}) + worker := startWorker(t, t.Context(), s, &execution.Dispatcher{}) t.Cleanup(func() { stopped, cancel := context.WithCancel(context.Background()) cancel() @@ -29,13 +29,13 @@ func TestInitialSessionInputOfficialClient(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, worker)) + handler, err := publicHandler(t, s, auth, "codex", workerExecution(t, worker)) if err != nil { t.Fatal(err) } server := httptest.NewServer(handler) defer server.Close() - unsupported, err := publicHandler(t, s, db, auth, "fake_alpha", workerExecution(t, worker)) + unsupported, err := publicHandler(t, s, auth, "fake_alpha", workerExecution(t, worker)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/session_model_execution_http_test.go b/services/core/tests/integration/session_model_execution_http_test.go index adb141a66..9a1e6c2f1 100644 --- a/services/core/tests/integration/session_model_execution_http_test.go +++ b/services/core/tests/integration/session_model_execution_http_test.go @@ -15,11 +15,10 @@ import ( func TestModelExecutionHTTPWriteOnlyAndStrictAdmission(t *testing.T) { _, pool := testStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{6}, 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: "catalog-test", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}}) - handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, db)) + handler, err := publicHandler(t, st, auth, "codex", storeExecution(t, st), managedSandboxes(t, st)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/session_model_options_test.go b/services/core/tests/integration/session_model_options_test.go index 9724b971e..553df763c 100644 --- a/services/core/tests/integration/session_model_options_test.go +++ b/services/core/tests/integration/session_model_options_test.go @@ -17,7 +17,7 @@ func TestSessionModelExecutionStoresOnlyProviderBundle(t *testing.T) { if err != nil { t.Fatal(err) } - st := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + st := NewWithCredentialCipher(pool, cipher) ctx, tenant := t.Context(), uuid.NewString() provider := &v1.ModelProviderInput{Protocol: "responses", BaseURL: "https://example.com/v1", APIKey: "provider-key-canary"} input := sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: []byte(`{"agent":{"model":"actual-model"},"environment":{"type":"openai_hosted"}}`), ModelProvider: provider} diff --git a/services/core/tests/integration/session_reads_fixture_test.go b/services/core/tests/integration/session_reads_fixture_test.go index 8ab6b3564..43f097c11 100644 --- a/services/core/tests/integration/session_reads_fixture_test.go +++ b/services/core/tests/integration/session_reads_fixture_test.go @@ -7,7 +7,6 @@ import ( "github.com/jackc/pgx/v5/pgxpool" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -16,8 +15,8 @@ import ( // tenant's hosted Environment on pool through the procedure the managed // Runtime allocation runs, without the allocation. func FixtureEnvironmentDevice(ctx context.Context, pool *pgxpool.Pool, tenant, environment, name, credentialHash string) (sessions.ExecutionDevice, error) { - units := pgunit.NewPool(pool) - current, err := sessionpg.New(units, nil).GetEnvironment(ctx, tenant, environment) + s := New(pool) + current, err := sessionAdapter(s).GetEnvironment(ctx, tenant, environment) if err != nil { return sessions.ExecutionDevice{}, err } @@ -30,7 +29,7 @@ func FixtureEnvironmentDevice(ctx context.Context, pool *pgxpool.Pool, tenant, e return sessions.ExecutionDevice{}, err } device := sessions.ExecutionDevice{ID: uuid.NewString(), Name: registration.Name, EnvironmentID: current.ID} - err = sessionpg.WithSession(ctx, units, lookup.TenantID, lookup.ID, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + err = sessionpg.WithSession(ctx, s.pooled, lookup.TenantID, lookup.ID, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { if err := locked.Public(); err != nil { return err } diff --git a/services/core/tests/integration/session_reference_retry_public_test.go b/services/core/tests/integration/session_reference_retry_public_test.go index 6e54d2c79..5ee70df0a 100644 --- a/services/core/tests/integration/session_reference_retry_public_test.go +++ b/services/core/tests/integration/session_reference_retry_public_test.go @@ -20,10 +20,10 @@ func TestSavedReferenceRetryOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, db := newTestStoreDB(t) + s, _ := testStore(t) tenant, token, foreign := uuid.NewString(), 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: tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}}) - worker := startWorker(t, t.Context(), db, &execution.Dispatcher{}) + worker := startWorker(t, t.Context(), s, &execution.Dispatcher{}) t.Cleanup(func() { ctx, cancel := context.WithCancel(context.Background()) cancel() @@ -31,13 +31,13 @@ func TestSavedReferenceRetryOfficialClient(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, worker)) + handler, err := publicHandler(t, s, auth, "codex", workerExecution(t, worker)) if err != nil { t.Fatal(err) } server := httptest.NewServer(handler) defer server.Close() - recovered, err := publicHandler(t, New(db.pool), db, auth, "codex") + recovered, err := publicHandler(t, New(s.pool), auth, "codex") if err != nil { t.Fatal(err) } @@ -56,9 +56,9 @@ func TestSavedReferenceRetryOfficialClient(t *testing.T) { } var err error if input.Delete { - _, err = db.pool.Exec(r.Context(), `DELETE FROM agents WHERE tenant_id=$1 AND id=$2`, tenant, input.ID) + _, err = s.pool.Exec(r.Context(), `DELETE FROM agents WHERE tenant_id=$1 AND id=$2`, tenant, input.ID) } else { - _, err = db.pool.Exec(r.Context(), `UPDATE agents SET configuration=configuration || $3::jsonb WHERE tenant_id=$1 AND id=$2`, tenant, input.ID, input.Patch) + _, err = s.pool.Exec(r.Context(), `UPDATE agents SET configuration=configuration || $3::jsonb WHERE tenant_id=$1 AND id=$2`, tenant, input.ID, input.Patch) } if err != nil { t.Error(err) diff --git a/services/core/tests/integration/sessions_fixture_test.go b/services/core/tests/integration/sessions_fixture_test.go deleted file mode 100644 index 0f921ab0a..000000000 --- a/services/core/tests/integration/sessions_fixture_test.go +++ /dev/null @@ -1,47 +0,0 @@ -package integration - -import ( - "testing" - - "github.com/jackc/pgx/v5/pgxpool" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -// sessionReads is the Session adapter on pool without a credential key. It -// serves the Session reads that open no frozen data. -func sessionReads(pool *pgxpool.Pool) *sessionpg.Store { - return sessionpg.New(pgunit.NewPool(pool), nil) -} - -// fixtureSessions builds the Session adapter and service on db, with the -// placement rules on db's public URL, as cmd/server does. -func fixtureSessions(db fixtureDB) (*sessionpg.Store, *sessions.Service, error) { - rules, err := placement.NewRules(providers.Builtin(), db.publicURL) - if err != nil { - return nil, nil, err - } - sessionStore := fixtureSessionStore(db) - sessionService, err := sessions.NewService(sessionStore, rules) - return sessionStore, sessionService, err -} - -// fixtureSessionStore is the Session adapter on db with its credential key, as -// cmd/server builds it. -func fixtureSessionStore(db fixtureDB) *sessionpg.Store { - return sessionpg.New(pgunit.NewPool(db.pool), db.cipher) -} - -// fixtureSessionService is the Session service of fixtureSessions(db). -func fixtureSessionService(t testing.TB, db fixtureDB) *sessions.Service { - t.Helper() - _, service, err := fixtureSessions(db) - if err != nil { - t.Fatal(err) - } - return service -} diff --git a/services/core/tests/integration/sessions_test.go b/services/core/tests/integration/sessions_test.go index ac8401aac..9e945672c 100644 --- a/services/core/tests/integration/sessions_test.go +++ b/services/core/tests/integration/sessions_test.go @@ -19,7 +19,7 @@ import ( func testStore(t *testing.T) (*Store, *pgxpool.Pool) { t.Helper() pool := pgtest.Open(t) - return withPlacement(t, New(pool)), pool + return New(pool), pool } // sessionAdapter is the Session adapter on s's database with s's credential diff --git a/services/core/tests/integration/skill_selectors_public_test.go b/services/core/tests/integration/skill_selectors_public_test.go index 42ccdcd51..7052c01b2 100644 --- a/services/core/tests/integration/skill_selectors_public_test.go +++ b/services/core/tests/integration/skill_selectors_public_test.go @@ -25,20 +25,20 @@ func TestSkillSelectorsOfficialClientPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) 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") + h, err := publicHandler(t, s, auth, "codex") if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() - recoveredStore, recoveredDB := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} - h, err = publicHandler(t, recoveredStore, recoveredDB, auth, "codex") + recoveredStore := NewWithCredentialCipher(pool, cipher) + h, err = publicHandler(t, recoveredStore, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/skill_version_deletion_public_test.go b/services/core/tests/integration/skill_version_deletion_public_test.go index d82d42a66..82ea71dce 100644 --- a/services/core/tests/integration/skill_version_deletion_public_test.go +++ b/services/core/tests/integration/skill_version_deletion_public_test.go @@ -25,13 +25,13 @@ func TestSkillVersionDeletionHTTPPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) owner, foreign, ownerTenant := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "skill-owner", TokenSHA256: runtimedevice.HashCredential(owner), TenantID: ownerTenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "skill-foreign", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - h, err := publicHandler(t, s, db, auth, "codex") + h, err := publicHandler(t, s, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/skill_version_deletion_test.go b/services/core/tests/integration/skill_version_deletion_test.go index 77547d92f..3a9b2e5a2 100644 --- a/services/core/tests/integration/skill_version_deletion_test.go +++ b/services/core/tests/integration/skill_version_deletion_test.go @@ -21,7 +21,7 @@ func TestSoleSkillVersionDeletionKeepsFrozenSetup(t *testing.T) { if err != nil { t.Fatal(err) } - s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) + s := NewWithCredentialCipher(pool, cipher) skillService := SkillService(t, pool, cipher) tenant := uuid.NewString() archive := skillArchive(t, "sole-version-frozen") diff --git a/services/core/tests/integration/skills_public_test.go b/services/core/tests/integration/skills_public_test.go index 5f1040746..4aa57368b 100644 --- a/services/core/tests/integration/skills_public_test.go +++ b/services/core/tests/integration/skills_public_test.go @@ -25,20 +25,20 @@ func TestSkillsOfficialClientPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) 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") + h, err := publicHandler(t, s, auth, "codex") if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() - recoveredStore, recoveredDB := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} - h, err = publicHandler(t, recoveredStore, recoveredDB, auth, "codex") + recoveredStore := NewWithCredentialCipher(pool, cipher) + h, err = publicHandler(t, recoveredStore, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/source_files_errors_public_test.go b/services/core/tests/integration/source_files_errors_public_test.go index 147d9ea3a..73fc3856b 100644 --- a/services/core/tests/integration/source_files_errors_public_test.go +++ b/services/core/tests/integration/source_files_errors_public_test.go @@ -17,13 +17,13 @@ func TestSourceFileErrorsOfficialClientPostgres(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: "files-owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: uuid.NewString()}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "files-foreign", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - h, err := publicHandler(t, s, db, auth, "codex") + h, err := publicHandler(t, s, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/stream_authority_http_test.go b/services/core/tests/integration/stream_authority_http_test.go index 97c287eeb..33dd4666a 100644 --- a/services/core/tests/integration/stream_authority_http_test.go +++ b/services/core/tests/integration/stream_authority_http_test.go @@ -22,8 +22,8 @@ func TestLiveStreamClosesAfterKeyRevocationOrProjectArchive(t *testing.T) { name = "project-archive" } t.Run(name, func(t *testing.T) { - s, db := newTestStoreDB(t) - _, management := fixtureProjects(t, db) + s, _ := testStore(t) + _, management := fixtureProjects(t, s) projectID := uuid.NewString() ctx := adminaudit.WithSource(t.Context(), adminaudit.Source{CredentialID: "12345678", ActorLabel: "test", RequestID: uuid.NewString(), TraceID: uuid.NewString(), ProjectID: projectID}) project, err := management.CreateProject(ctx, projects.CreateProject{ID: projectID, Name: "Stream authority"}) @@ -42,7 +42,7 @@ func TestLiveStreamClosesAfterKeyRevocationOrProjectArchive(t *testing.T) { if err != nil { t.Fatal(err) } - h, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), storeExecution(t, s)) + h, err := publicHandler(t, s, nil, "codex", storeKeys(s), storeExecution(t, s)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/structured_output_dispatch_test.go b/services/core/tests/integration/structured_output_dispatch_test.go index ed70760cf..2ae4b27be 100644 --- a/services/core/tests/integration/structured_output_dispatch_test.go +++ b/services/core/tests/integration/structured_output_dispatch_test.go @@ -27,7 +27,7 @@ func TestStructuredOutputDispatchRechecksOperationQualification(t *testing.T) { if err != nil { t.Fatal(err) } - 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) } caps := prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, StructuredOutput: proto.CapabilitySupported, MessageItems: proto.CapabilitySupported, Preparation: proto.CapabilitySupported}) diff --git a/services/core/tests/integration/structured_output_native_test.go b/services/core/tests/integration/structured_output_native_test.go index c1abc64e6..afa3de2b6 100644 --- a/services/core/tests/integration/structured_output_native_test.go +++ b/services/core/tests/integration/structured_output_native_test.go @@ -28,7 +28,7 @@ func TestNativeStructuredOutputPublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) 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() { @@ -44,7 +44,7 @@ func TestNativeStructuredOutputPublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, "claude_sdk", workerExecution(t, worker), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, auth, "claude_sdk", workerExecution(t, worker), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } @@ -69,7 +69,7 @@ func TestNativeStructuredOutputPublicExecution(t *testing.T) { if err != nil || json.Unmarshal(raw, &proof) != nil { t.Fatal("invalid evidence", err) } - call, err := FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turn, proof.Call) + call, err := FixtureFunctionCall(ctx, h.s.pool, h.tenant, proof.Session, proof.Turn, proof.Call) if err != nil || !call.Applied { t.Fatal("function application receipt missing", err) } diff --git a/services/core/tests/integration/subagent_dispatch_test.go b/services/core/tests/integration/subagent_dispatch_test.go index 6a937f75e..e56fa1df2 100644 --- a/services/core/tests/integration/subagent_dispatch_test.go +++ b/services/core/tests/integration/subagent_dispatch_test.go @@ -23,7 +23,7 @@ func TestSubagentIdentityUsesLeasedDispatchJournal(t *testing.T) { if err != nil { t.Fatal(err) } - 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) } input := h.message("first", "root message") diff --git a/services/core/tests/integration/subagent_visibility_public_test.go b/services/core/tests/integration/subagent_visibility_public_test.go index 7a71affb0..b6b7ec69e 100644 --- a/services/core/tests/integration/subagent_visibility_public_test.go +++ b/services/core/tests/integration/subagent_visibility_public_test.go @@ -70,20 +70,20 @@ func subagentFixture(kind string, value any) sessions.ExecutionEvent { // routes with the Session's Agent ID, Subagent lists use the common envelope and // child Item lists clamp their limit. Tenant B sees none of it. func TestSubagentVisibilityPublic(t *testing.T) { - s, db := newTestStoreDB(t) + s, _ := testStore(t) tenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "foreign", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) + handler, err := publicHandler(t, s, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } server := httptest.NewServer(handler) defer server.Close() client := pathIDClient{t: t, server: server} - leased := executionOwner(t, db) + leased := executionOwner(t, s) writer := NewExecution(s, leased.Lease.(*pgunit.Lease)) ctx := t.Context() @@ -121,7 +121,7 @@ func TestSubagentVisibilityPublic(t *testing.T) { t.Fatal(page, err) } root := page.Turns[0].ID - host, err := fixtureSessionService(t, db).CreateDevice(ctx, tenant, "subagent visibility", runtimedevice.HashCredential(uuid.NewString())) + host, err := sessionService(t, s).CreateDevice(ctx, tenant, "subagent visibility", runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/template_composition_public_test.go b/services/core/tests/integration/template_composition_public_test.go index 76eaa9e2f..a958e47ac 100644 --- a/services/core/tests/integration/template_composition_public_test.go +++ b/services/core/tests/integration/template_composition_public_test.go @@ -31,9 +31,6 @@ func TestTemplateCompositionOfficialClientPostgres(t *testing.T) { t.Fatal(err) } s, reopenedStore := NewWithCredentialCipher(pool, cipher), NewWithCredentialCipher(pool, cipher) - db := fixtureDB{pool: pool, cipher: cipher} // built both Stores - s.SetPlacement(fixtureRules(t, db)) - reopenedStore.SetPlacement(fixtureRules(t, db)) tenant, foreignTenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "composition-owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, @@ -42,7 +39,7 @@ func TestTemplateCompositionOfficialClientPostgres(t *testing.T) { serve := func(current *Store) *httptest.Server { t.Helper() // Hosted admission and freezing use the real Store; no Runtime or model runs. - h, err := publicHandler(t, current, db, auth, "codex", storeExecution(t, current), managedSandboxes(t, db), fixtureDeploymentProvider()) + h, err := publicHandler(t, current, auth, "codex", storeExecution(t, current), managedSandboxes(t, current), fixtureDeploymentProvider()) if err != nil { t.Fatal(err) } @@ -105,7 +102,7 @@ func TestTemplateCompositionOfficialClientPostgres(t *testing.T) { t.Fatalf("unknown case %q", label) } // A reader built after the requests reads the frozen setup and files. - current := fixtureSessionStore(db) + current := sessionAdapter(s) setup, err := current.ReadEnvironmentSetup(t.Context(), tenant, id) if err != nil || !reflect.DeepEqual(setup.Env, env) || !reflect.DeepEqual(setup.PackageMetadata(), packages) || len(setup.Commands) != len(commands) { t.Fatalf("%s durable setup differs: %v", label, err) diff --git a/services/core/tests/integration/template_null_selection_public_test.go b/services/core/tests/integration/template_null_selection_public_test.go index 294166b91..368e7f5e5 100644 --- a/services/core/tests/integration/template_null_selection_public_test.go +++ b/services/core/tests/integration/template_null_selection_public_test.go @@ -33,9 +33,6 @@ func TestTemplateNullSelectionOfficialClientPostgres(t *testing.T) { t.Fatal(err) } s, reopenedStore := NewWithCredentialCipher(pool, cipher), NewWithCredentialCipher(pool, cipher) - db := fixtureDB{pool: pool, cipher: cipher} // built both Stores - s.SetPlacement(fixtureRules(t, db)) - reopenedStore.SetPlacement(fixtureRules(t, db)) tenant, foreignTenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "selection-owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, @@ -43,7 +40,7 @@ func TestTemplateNullSelectionOfficialClientPostgres(t *testing.T) { }) serve := func(current *Store) *httptest.Server { t.Helper() - h, err := publicHandler(t, current, db, auth, "codex", storeExecution(t, current), managedSandboxes(t, db), fixtureDeploymentProvider()) + h, err := publicHandler(t, current, auth, "codex", storeExecution(t, current), managedSandboxes(t, current), fixtureDeploymentProvider()) if err != nil { t.Fatal(err) } @@ -95,7 +92,7 @@ func TestTemplateNullSelectionOfficialClientPostgres(t *testing.T) { t.Fatalf("missing expectation for %s", label) } // A reader built after the requests reads the frozen setup and files. - current := fixtureSessionStore(db) + current := sessionAdapter(s) setup, err := current.ReadEnvironmentSetup(t.Context(), tenant, id) if err != nil { t.Fatalf("%s frozen setup: %v", label, err) diff --git a/services/core/tests/integration/token_usage_integration_test.go b/services/core/tests/integration/token_usage_integration_test.go index de9dba494..6b533dc6b 100644 --- a/services/core/tests/integration/token_usage_integration_test.go +++ b/services/core/tests/integration/token_usage_integration_test.go @@ -16,7 +16,7 @@ import ( func TestTokenUsageDurableSnapshotsAndSessionTotals(t *testing.T) { ctx := context.Background() s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "usage"}) if err != nil { @@ -106,8 +106,8 @@ func TestTokenUsageDurableSnapshotsAndSessionTotals(t *testing.T) { func TestCancellationReceiptUsageSurvivesRecovery(t *testing.T) { ctx := context.Background() - s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + s, _ := testStore(t) + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "cancel-recovery"}) if err != nil { @@ -141,8 +141,8 @@ func TestCancellationReceiptUsageSurvivesRecovery(t *testing.T) { // and after a Turn ends with unknown usage (EVT-13). func TestSessionUsageRequiresEveryRootTurnEndedAndMeasured(t *testing.T) { ctx := context.Background() - s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + s, _ := testStore(t) + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "unknown-usage"}) if err != nil { diff --git a/services/core/tests/integration/tool_policy_native_test.go b/services/core/tests/integration/tool_policy_native_test.go index 4b443aee4..8f5030ca3 100644 --- a/services/core/tests/integration/tool_policy_native_test.go +++ b/services/core/tests/integration/tool_policy_native_test.go @@ -34,7 +34,7 @@ func TestNativeToolPolicyPublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 15*time.Minute) 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() { @@ -50,7 +50,7 @@ func TestNativeToolPolicyPublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: foreignTenant, SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: foreignTenant}, }) - handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(t, worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, auth, kind, workerExecution(t, worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/tool_search_native_test.go b/services/core/tests/integration/tool_search_native_test.go index 2312b1c59..f5e950500 100644 --- a/services/core/tests/integration/tool_search_native_test.go +++ b/services/core/tests/integration/tool_search_native_test.go @@ -28,7 +28,7 @@ func TestNativeToolSearchPublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) 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() { @@ -44,7 +44,7 @@ func TestNativeToolSearchPublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, "claude_sdk", workerExecution(t, worker), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, auth, "claude_sdk", workerExecution(t, worker), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } @@ -69,7 +69,7 @@ func TestNativeToolSearchPublicExecution(t *testing.T) { if err != nil || json.Unmarshal(raw, &proof) != nil { t.Fatal("invalid evidence", err) } - call, err := FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turn, proof.Call) + call, err := FixtureFunctionCall(ctx, h.s.pool, h.tenant, proof.Session, proof.Turn, proof.Call) if err != nil || !call.Applied { t.Fatal("function application receipt missing", err) } diff --git a/services/core/tests/integration/turn_events_test.go b/services/core/tests/integration/turn_events_test.go index 3c5a39839..e95bf1a40 100644 --- a/services/core/tests/integration/turn_events_test.go +++ b/services/core/tests/integration/turn_events_test.go @@ -14,7 +14,7 @@ import ( func TestTurnEventBatchesAreOrderedIsolatedAndDurable(t *testing.T) { ctx := context.Background() s, pool := testStore(t) - journal := executionOwner(t, fixtureDB{pool: pool}).Sessions + journal := executionOwner(t, s).Sessions tenant := uuid.NewString() session, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "events"}) if err != nil { diff --git a/services/core/tests/integration/unified_model_configuration_http_test.go b/services/core/tests/integration/unified_model_configuration_http_test.go index 769a77d21..bfcea9301 100644 --- a/services/core/tests/integration/unified_model_configuration_http_test.go +++ b/services/core/tests/integration/unified_model_configuration_http_test.go @@ -14,14 +14,14 @@ import ( ) func TestUnifiedModelConfigurationHTTP(t *testing.T) { - st, db := newManagedTestStoreDB(t) + st, _ := newManagedTestStore(t) tenant, token, coreKey := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "model-configuration", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}}) admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential(coreKey)}) if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, db), withCoreKeys(admin), withHarnesses([]string{"codex", "claude_sdk"})) + handler, err := publicHandler(t, st, auth, "codex", storeExecution(t, st), managedSandboxes(t, st), withCoreKeys(admin), withHarnesses([]string{"codex", "claude_sdk"})) if err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/unstorable_text_public_test.go b/services/core/tests/integration/unstorable_text_public_test.go index fec595786..01f1f9804 100644 --- a/services/core/tests/integration/unstorable_text_public_test.go +++ b/services/core/tests/integration/unstorable_text_public_test.go @@ -24,10 +24,10 @@ func TestUnstorableTextRejectsWithoutWritesPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s, db := NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s := NewWithCredentialCipher(pool, cipher) token := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "nul-owner", TokenSHA256: runtimedevice.HashCredential(token), 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) } diff --git a/services/core/tests/integration/vaults_fixture_test.go b/services/core/tests/integration/vaults_fixture_test.go index d5d8697fd..5249ca3ea 100644 --- a/services/core/tests/integration/vaults_fixture_test.go +++ b/services/core/tests/integration/vaults_fixture_test.go @@ -7,14 +7,14 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/vaults" ) -// fixtureVaults builds the Vault adapter and service on db, as cmd/server does. -// A keyless db leaves the operations that need no credential key available. -func fixtureVaults(db fixtureDB) (*vaultpg.Store, *vaults.Service, error) { +// fixtureVaults builds the Vault adapter and service on s, as cmd/server does. +// A keyless s leaves the operations that need no credential key available. +func fixtureVaults(s *Store) (*vaultpg.Store, *vaults.Service, error) { refresher, err := oauthrefresh.NewClient(nil) if err != nil { return nil, nil, err } - vaultStore := vaultpg.New(pgunit.NewPool(db.pool), db.cipher) + vaultStore := vaultpg.New(pgunit.NewPool(s.pool), s.credentialCipher) vaultService, err := vaults.NewService(vaultStore, refresher) return vaultStore, vaultService, err } diff --git a/services/core/tests/integration/whitespace_input_public_test.go b/services/core/tests/integration/whitespace_input_public_test.go index 5bfeeafec..3954c0b49 100644 --- a/services/core/tests/integration/whitespace_input_public_test.go +++ b/services/core/tests/integration/whitespace_input_public_test.go @@ -19,10 +19,10 @@ import ( // still reject with today's fields and write nothing. func TestWhitespaceInputStoredVerbatimPostgres(t *testing.T) { // An isolated database keeps the no-write digest independent of other tests. - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) token := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "whitespace-owner", TokenSHA256: runtimedevice.HashCredential(token), 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) } @@ -72,7 +72,7 @@ func TestWhitespaceInputStoredVerbatimPostgres(t *testing.T) { } // W4: unchanged rejection without writes. - before := databaseDigest(t, db.pool) + before := databaseDigest(t, s.pool) const rejection = `{"error":{"message":"Invalid resource identifier or request limits.","type":"invalid_request_error","code":"invalid_request","param":null}}` + "\n" for _, input := range []string{`""`, `[]`, `[{"role":"user","content":[]}]`, `[{"role":"user","content":[{"type":"input_text","text":""}]}]`} { if status, body := client.do(token, http.MethodPost, "/v1/agents/sessions", "application/json", []byte(`{"agent":{"model":"whitespace-model"},"environment":{"type":"none"},"input":`+input+`}`)); status != http.StatusBadRequest || body != rejection { @@ -84,7 +84,7 @@ func TestWhitespaceInputStoredVerbatimPostgres(t *testing.T) { t.Errorf("events %s: %d %s", input, status, body) } } - if after := databaseDigest(t, db.pool); !mapsEqual(before, after) { + if after := databaseDigest(t, s.pool); !mapsEqual(before, after) { t.Error("rejected empty input changed persisted state") } } @@ -93,11 +93,11 @@ func TestWhitespaceInputStoredVerbatimPostgres(t *testing.T) { // admits it; Claude SDK and MiniMax Code reject it at Session creation and // events.create, before any write, reservation or promotion. func TestWhitespaceOnlyTextHarnessAdmissionPostgres(t *testing.T) { - s, db := newManagedTestStoreDB(t) + s, _ := newManagedTestStore(t) token := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "whitespace-harness", TokenSHA256: runtimedevice.HashCredential(token), TenantID: uuid.NewString()}}) // Real Worker admission with dispatch paused keeps admitted Turns queued. - worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Registry: runtimegateway.NewRegistry()}) + worker := startWorker(t, t.Context(), s, &execution.Dispatcher{Registry: runtimegateway.NewRegistry()}) t.Cleanup(func() { stopped, cancel := context.WithCancel(context.Background()) cancel() @@ -106,7 +106,7 @@ func TestWhitespaceOnlyTextHarnessAdmissionPostgres(t *testing.T) { } }) serve := func(engine string) pathIDClient { - handler, err := publicHandler(t, s, db, auth, engine, workerExecution(t, worker), executorURL("https://offline-executor.example")) + handler, err := publicHandler(t, s, auth, engine, workerExecution(t, worker), executorURL("https://offline-executor.example")) if err != nil { t.Fatal(err) } @@ -139,7 +139,7 @@ func TestWhitespaceOnlyTextHarnessAdmissionPostgres(t *testing.T) { if status, body := events(client, session, cancel); status != http.StatusAccepted { t.Fatalf("%s cancel: %d %s", engine, status, body) } - before := databaseDigest(t, db.pool) + before := databaseDigest(t, s.pool) for _, body := range []string{ `{"agent":{"model":"m"},"environment":{"type":"none"},"input":" "}`, `{"agent":{"model":"m"},"environment":{"type":"none"},"input":[{"role":"user","content":[{"type":"input_text","text":"\n\t"}]}]}`, @@ -157,7 +157,7 @@ func TestWhitespaceOnlyTextHarnessAdmissionPostgres(t *testing.T) { t.Errorf("%s events %s: %d %s", engine, body, status, response) } } - if after := databaseDigest(t, db.pool); !mapsEqual(before, after) { + if after := databaseDigest(t, s.pool); !mapsEqual(before, after) { t.Errorf("%s: rejected whitespace-only text changed persisted state", engine) } // Whitespace beside non-whitespace text in one message remains admitted verbatim. diff --git a/services/core/tests/integration/worker_capacity_test.go b/services/core/tests/integration/worker_capacity_test.go index f125e435c..0f11cef48 100644 --- a/services/core/tests/integration/worker_capacity_test.go +++ b/services/core/tests/integration/worker_capacity_test.go @@ -42,7 +42,7 @@ func TestWorkerDefersPreparationCapacityUntilCleanupReleasesSlot(t *testing.T) { turns[h.session.ID] = receipt.TurnID } frames := capacityWorkerFrames(t, h) - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) defer stop() admissions := make(map[string]string) started := make(map[string]int) @@ -175,7 +175,7 @@ func TestWorkerDoesNotDeferOtherPreparationOrStartRejections(t *testing.T) { h.session = publicSession(t, h, test.name) receipt := h.message("work", "once") frames := capacityWorkerFrames(t, h) - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) defer stop() prepare := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) if test.operation == proto.TypeExecutionStart { diff --git a/services/core/tests/integration/worker_fixture_test.go b/services/core/tests/integration/worker_fixture_test.go new file mode 100644 index 000000000..872c77345 --- /dev/null +++ b/services/core/tests/integration/worker_fixture_test.go @@ -0,0 +1,112 @@ +package integration + +import ( + "context" + "errors" + "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/modelconfigurationpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// startWorker starts the execution Worker as cmd/server does: it acquires the +// execution lease on s's database and hands it, with the execution writer built +// on it, to the Worker, which closes it when Run exits. The Worker opens MCP +// bearer tokens through the vaults service on s, records model configuration +// observations through the model configuration adapter on s, runs Session use +// cases and reads through the Session service and adapter on s, and reads the +// deployment through the deployment adapter on s. +func startWorker(t testing.TB, ctx context.Context, s *Store, dispatcher *execution.Dispatcher) *execution.Worker { + t.Helper() + worker, err := startWorkerErr(ctx, s, dispatcher) + if err != nil { + t.Fatal(err) + } + return worker +} + +// startWorkerErr is startWorker for tests that assert a startup failure. +func startWorkerErr(ctx context.Context, s *Store, dispatcher *execution.Dispatcher) (*execution.Worker, error) { + lease, err := pgunit.AcquireLease(ctx, s.pool) + if err != nil { + return nil, err + } + owner, err := fixtureOwner(s, lease) + if err != nil { + return nil, errors.Join(err, lease.Close(ctx)) + } + return startOwnedWorkerErr(ctx, s, dispatcher, owner) +} + +// startOwnedWorker is startWorker on an Owner the test already holds, for tests +// that also run execution operations on it. The Worker closes its lease when +// Run exits. +func startOwnedWorker(t testing.TB, ctx context.Context, s *Store, dispatcher *execution.Dispatcher, owner execution.Owner) *execution.Worker { + t.Helper() + worker, err := startOwnedWorkerErr(ctx, s, dispatcher, owner) + if err != nil { + t.Fatal(err) + } + return worker +} + +func startOwnedWorkerErr(ctx context.Context, s *Store, dispatcher *execution.Dispatcher, owner execution.Owner) (*execution.Worker, error) { + _, credentials, err := fixtureVaults(s) + if err != nil { + return nil, errors.Join(err, owner.Lease.Close(ctx)) + } + service, err := newSessionService(s) + if err != nil { + return nil, errors.Join(err, owner.Lease.Close(ctx)) + } + deployments, err := fixtureDeploymentService(s) + if err != nil { + return nil, errors.Join(err, owner.Lease.Close(ctx)) + } + owned := *dispatcher + owned.Credentials = credentials + owned.Observer = modelconfigurationpg.New(pgunit.NewPool(s.pool), s.credentialCipher) + owned.Deployment = deployments + owned.DeploymentReader = deploymentStore(s) + owned.Sessions = service + owned.SessionsReader = sessionAdapter(s) + return execution.StartWorker(ctx, &owned, owner) +} + +// executionOwner acquires the execution lease on s's database and builds the +// execution operations on it, for tests that run them without a Worker. The +// lease closes when the test ends. +func executionOwner(t testing.TB, s *Store) execution.Owner { + t.Helper() + lease, err := pgunit.AcquireLease(t.Context(), s.pool) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = lease.Close(context.Background()) }) + owner, err := fixtureOwner(s, lease) + if err != nil { + t.Fatal(err) + } + return owner +} + +// fixtureOwner builds the deployment and Session execution operations on +// lease, as cmd/server does. +func fixtureOwner(s *Store, lease *pgunit.Lease) (execution.Owner, error) { + _, changes, err := fixtureDeploymentExecution(s, lease) + if err != nil { + return execution.Owner{}, err + } + sessionExecution, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) + if err != nil { + return execution.Owner{}, err + } + return execution.Owner{ + Lease: lease, + Deployment: changes, + Sessions: sessionExecution, + }, nil +} diff --git a/services/core/tests/integration/worker_input_race_test.go b/services/core/tests/integration/worker_input_race_test.go index eaa856fbb..76030a770 100644 --- a/services/core/tests/integration/worker_input_race_test.go +++ b/services/core/tests/integration/worker_input_race_test.go @@ -59,7 +59,7 @@ func TestWorkerInputReadSkipsConcurrentlyCancelledCandidate(t *testing.T) { defer instrumented.Close() ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker := startWorker(t, ctx, fixtureDB{pool: instrumented}, h.d) + worker := startWorker(t, ctx, New(instrumented), h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { diff --git a/services/core/tests/integration/worker_lease_loss_test.go b/services/core/tests/integration/worker_lease_loss_test.go index bedf0f7c7..d2b75293f 100644 --- a/services/core/tests/integration/worker_lease_loss_test.go +++ b/services/core/tests/integration/worker_lease_loss_test.go @@ -17,7 +17,7 @@ func TestWorkerLeaseLossLeavesUncertainWorkForSuccessor(t *testing.T) { h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) h.session = publicSession(t, h, "active") queued := publicSession(t, h, "queued") - worker := startWorker(t, t.Context(), h.db, h.d) + worker := startWorker(t, t.Context(), h.s, h.d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -71,7 +71,7 @@ func TestWorkerLeaseLossLeavesUncertainWorkForSuccessor(t *testing.T) { if err != nil || active.Status != sessions.TurnInProgress { t.Fatal("lost owner persisted fallback completion", active, err) } - successor := startWorker(t, t.Context(), h.db, h.d) + successor := startWorker(t, t.Context(), h.s, h.d) stopped, stop := context.WithCancel(t.Context()) stop() if err = successor.Run(stopped); err != context.Canceled { diff --git a/services/core/tests/integration/worker_preparation_failure_test.go b/services/core/tests/integration/worker_preparation_failure_test.go index 4a744e311..d1d266352 100644 --- a/services/core/tests/integration/worker_preparation_failure_test.go +++ b/services/core/tests/integration/worker_preparation_failure_test.go @@ -19,7 +19,7 @@ func TestWorkerSettlesConfirmedPreparationFailureAndAcceptsNewInput(t *testing.T if err != nil { t.Fatal(err) } - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) defer stop() prepare := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) if code == "preparation_failed" { @@ -102,7 +102,7 @@ func TestWorkerRetriesUncertainPreparationFailure(t *testing.T) { if err != nil { t.Fatal(err) } - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) defer stop() prepare := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) status := proto.PreparationStatusPayload{State: response.state, Operation: response.operation, ErrorCode: response.code, RunID: response.runID} @@ -134,7 +134,7 @@ func TestWorkerPreparationRejectionPreservesCancellationAndNewerInput(t *testing if err != nil { t.Fatal(err) } - _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.s, h.d) defer stop() old := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) if _, err := cancelEnvironmentInput(t.Context(), h.s, h.tenant, h.session.ID, first.ID); err != nil { diff --git a/services/core/tests/integration/worker_wakeup_test.go b/services/core/tests/integration/worker_wakeup_test.go index 379cd2349..3596bbff1 100644 --- a/services/core/tests/integration/worker_wakeup_test.go +++ b/services/core/tests/integration/worker_wakeup_test.go @@ -46,7 +46,7 @@ func TestWorkerSchedulerCommittedAdmissionWakesBeforeMaintenance(t *testing.T) { defer instrumented.Close() ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker := startWorker(t, ctx, fixtureDB{pool: instrumented, cipher: fixtureCipher}, h.d) + worker := startWorker(t, ctx, NewWithCredentialCipher(instrumented, fixtureCipher), h.d) done := make(chan error, 1) started := false defer func() { @@ -102,7 +102,7 @@ func TestWorkerSchedulerHintBypassesEnvironmentScanThrottle(t *testing.T) { h := newDispatchHarnessForSession(t, []byte(`{"agent":{"model":"test-model"},"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`), false) enableWorkerEnvironment(t, h) frames := workerFrames(t, h) - worker, stop := startEnvironmentExpiryWorker(t, h.db, h.d) + worker, stop := startEnvironmentExpiryWorker(t, h.s, h.d) defer stop() awaitDaemonRemoteCondition(t, t.Context(), 5*time.Second, "initial empty scheduler scan", func() bool { return worker.MetricsSnapshot().Scheduler.LastRunAt != nil