diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 5dd35520..1d6fe441 100644 --- a/services/core/IMPLEMENTATION.md +++ b/services/core/IMPLEMENTATION.md @@ -6,11 +6,13 @@ These are the code-level rules of `services/core` that no contract states. Contr `api.NewHandler` takes one `api.Dependencies` value, built only in `cmd/server`. Each application area is one field typed as an interface declared in `api` beside its handlers, listing exactly the methods they call. Every field is required and `NewHandler` rejects a missing one, except the optional groups whose comments say what nil means: `Execution` is nil without an execution Worker, `Sandboxes` is nil without a managed sandbox installation and requires `Execution`, and `Execution.NativeInstaller` is nil for a build without a source revision. Handlers never discover a capability by type assertion or fall back to another implementation. API tests use one strict fake per area, `fake`, which fails the test on any call the test did not set. -`internal/persistence/postgres/pgunit` owns Core's PostgreSQL transaction and execution-lease mechanics: pooled read-write and snapshot transactions, the lease's dedicated connection and its gate, the ownership check, the cancellation fence, close, and the execution deadline. Persistence code runs every transaction through it, and nothing outside `persistence` and `store` imports it. `internal/persistence/postgres/pgtest` is test support: it opens the dedicated test database under the `oac_*_tests` guard, applies the migrations, and creates isolated databases for database-wide state such as the execution lease. Only test files import it. +`internal/persistence/postgres/pgunit` owns Core's PostgreSQL transaction and execution-lease mechanics: pooled read-write and snapshot transactions, the lease's dedicated connection and its gate, the ownership check, the cancellation fence, close, and the execution deadline. Persistence code runs every transaction through it. Outside `persistence` and `store`, only `cmd/server`, which acquires the lease, and test fixtures import it. `internal/persistence/postgres/pgtest` is test support: it opens the dedicated test database under the `oac_*_tests` guard, applies the migrations, and creates isolated databases for database-wide state such as the execution lease. Only test files import it. Shared vocabulary has one owner each, and domains use it rather than copy it. `internal/environmentconfig` owns Environment setup, Skills, Plugins and initial files with their validation and public metadata; `Setup.Validate` checks requested configuration, where a Skill may be an unresolved reference, and `Setup.ValidateInstalled` checks frozen, installable configuration. `internal/skills` owns `ParseVersion`, the canonical positive decimal Skill version. `internal/metadata` owns the metadata rules: `Validate` for the pair, key and value limits and U+0000, `ValidateStorable` for U+0000 alone, and `Encode` with its 64 KiB bound. `internal/jsonobject` owns `Normalize`, the stable encoding of stored JSON objects that snapshots and retry identities compare. These packages import no persistence. -`store` is transitional. `store.New` builds a pooled Store, and `store.NewExecution` takes the lease and builds the execution writer on it. An execution-only operation on a pooled Store fails with `store.ErrExecutionAuthority`. New adapters do not copy that check: their execution repositories require a `*pgunit.Lease` at construction, their public repositories expose no execution operation, and the check goes away with `store`. +`store` is transitional. `store.New` builds a pooled Store, and `store.NewExecution(s, lease)` builds the execution writer on a lease it borrows. An execution-only operation on a pooled Store fails with `store.ErrExecutionAuthority`. New adapters do not copy that check: their execution repositories require a `*pgunit.Lease` at construction, their public repositories expose no execution operation, and the check goes away with `store`. + +`cmd/server` owns the execution lease. It acquires one `pgunit.Lease`, builds every lease-bound adapter on it, and passes the lease and those adapters together as one `execution.Owner` to `execution.StartWorker`. If anything fails before that call, `cmd/server` closes the lease. From that call the Worker owns cleanup: a failed start closes the lease before it returns, and a started Worker closes it after `Run` has cancelled and drained its work. Each close runs under its own bounded deadline, independent of the cancelled request or run. Lease-bound adapters and the store writer borrow the lease and never close it, and the Worker uses the lease only through `Owner.Lease`, never through an adapter. Store integration tests start the Worker the same way through `startWorker`. ## Request handling @@ -153,7 +155,7 @@ Item merging never mutates the incoming observation or the previous snapshot: pu ## Worker ownership -Enabling the daemon gateway with `OAC_PUBLIC_URL` also starts the execution Worker. One Worker owns an execution database through a `pgunit.Lease`: the PostgreSQL advisory lock held by one dedicated connection. A second Worker on the same database fails to start. The Worker's execution writer runs every Session transaction on that connection: binding, claim and reconciliation, journal, Items and usage, function callbacks and receipts, and terminal state. The lease gate serializes these short transactions and the ownership pings, and `pgunit.ExecutionTimeout` (five seconds) bounds each one, including its gate and Session-lock waits. Cancelling an in-flight pgx operation can close the connection that owns the lock, so coordinator-owned work is cancelled only through the lease's cancellation fence, between leased operations. Never hold a transaction across daemon or model work, reconnect the writer or fall back to the pool after losing the lease. Public admission and device maintenance use pooled connections, and pooled reads grant no write authority. Transactions state their isolation: read committed for writes, read-only repeatable read for snapshots. +Enabling the daemon gateway with `OAC_PUBLIC_URL` also starts the execution Worker. One Worker owns an execution database through a `pgunit.Lease`: the PostgreSQL advisory lock held by one dedicated connection. A second Core on the same database cannot acquire the lease and starts no Worker. The Worker's execution writer runs every Session transaction on that connection: binding, claim and reconciliation, journal, Items and usage, function callbacks and receipts, and terminal state. The lease gate serializes these short transactions and the ownership pings, and `pgunit.ExecutionTimeout` (five seconds) bounds each one, including its gate and Session-lock waits. Cancelling an in-flight pgx operation can close the connection that owns the lock, so coordinator-owned work is cancelled only through the lease's cancellation fence, between leased operations. Never hold a transaction across daemon or model work, reconnect the writer or fall back to the pool after losing the lease. Public admission and device maintenance use pooled connections, and pooled reads grant no write authority. Transactions state their isolation: read committed for writes, read-only repeatable read for snapshots. Sandbox reset snapshots bind the deployment relation explicitly to its single row before joining resources, so that even on a fresh database without statistics an inflated join estimate cannot trigger JIT compilation inside the lease deadline. diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index 79b0c20e..20cf8f05 100644 --- a/services/core/cmd/server/main.go +++ b/services/core/cmd/server/main.go @@ -36,6 +36,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/databaseurl" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/nativeinstaller" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtime" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeenrollment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" @@ -226,8 +227,12 @@ func run() error { if registry != nil { dispatcher := &execution.Dispatcher{Store: executionStore, Registry: registry, ManagedRuntimes: managed, MaxConcurrentExecutions: concurrency} - - worker, err = execution.StartWorker(ctx, dispatcher) + lease, err := pgunit.AcquireLease(ctx, pool) + if err != nil { + return err + } + // From this call on the Worker closes the lease, even when it fails to start. + worker, err = execution.StartWorker(ctx, dispatcher, execution.Owner{Lease: lease, Store: store.NewExecution(executionStore, lease)}) if err != nil { return err } diff --git a/services/core/internal/execution/archive_cancellation_cleanup_test.go b/services/core/internal/execution/archive_cancellation_cleanup_test.go index b8759fae..f02edbf7 100644 --- a/services/core/internal/execution/archive_cancellation_cleanup_test.go +++ b/services/core/internal/execution/archive_cancellation_cleanup_test.go @@ -62,7 +62,8 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { }{{"Kill_no_delivery", false, false}, {"KillCompute_no_delivery", true, false}, {"Kill_live_delivery", false, true}, {"KillCompute_live_delivery", true, true}} { t.Run(scenario.name, func(t *testing.T) { checkpoint := scenario.checkpoint - s, writer := resetManagerStore(t) + s, leased := resetManagerStore(t) + writer := leased.Store installation := uuid.NewString() if err := writer.ClaimWebSandboxDeployment(t.Context(), installation); err != nil { t.Fatal(err) @@ -167,7 +168,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { t.Fatal("Kill bypassed durable cleanup ownership", allocation, err) } }} - lifecycle := &runtimeLifecycle{store: writer, registry: registry, config: RuntimeProvider{InstallationID: installation, Provider: provider}, connections: map[string]*runtimeConnection{}} + lifecycle := &runtimeLifecycle{store: writer, lease: leased.Lease, registry: registry, config: RuntimeProvider{InstallationID: installation, Provider: provider}, connections: map[string]*runtimeConnection{}} if checkpoint { lifecycle.config.Provider = waitingCleanupCheckpoint{beforeKill: provider.beforeKill} } diff --git a/services/core/internal/execution/deployment_provider_observations_test.go b/services/core/internal/execution/deployment_provider_observations_test.go index e223208e..97f7fbe8 100644 --- a/services/core/internal/execution/deployment_provider_observations_test.go +++ b/services/core/internal/execution/deployment_provider_observations_test.go @@ -19,6 +19,7 @@ import ( type finishObservationFixture struct { s *store.Store writer *store.Store + lease Ownership pool *pgxpool.Pool tenant string session store.Session @@ -28,7 +29,7 @@ type finishObservationFixture struct { func newFinishObservationFixture(t *testing.T, maxConnections int32) finishObservationFixture { t.Helper() var cfg *pgxpool.Config - s, writer := resetManagerStoreConfig(t, func(c *pgxpool.Config) { + s, owner := resetManagerStoreConfig(t, func(c *pgxpool.Config) { if maxConnections > 0 { c.MaxConns = maxConnections } @@ -56,7 +57,7 @@ func newFinishObservationFixture(t *testing.T, maxConnections int32) finishObser if err != nil { t.Fatal(err) } - return finishObservationFixture{s, writer, pool, tenant, session, Dispatcher{Store: writer}} + return finishObservationFixture{s, owner.Store, owner.Lease, pool, tenant, session, Dispatcher{Store: owner.Store}} } func (f finishObservationFixture) start(t *testing.T) store.InputReceipt { t.Helper() @@ -183,7 +184,7 @@ func TestFinishRunObservationLockTimeoutAndFailureKeepLease(t *testing.T) { if elapsed := time.Since(started); elapsed < 900*time.Millisecond || elapsed > 2*time.Second { t.Fatal("pool timeout exceeded observation budget", elapsed) } - if err = f.writer.CheckExecutionOwnership(t.Context()); err != nil { + if err = f.lease.CheckOwnership(t.Context()); err != nil { t.Fatal("observation cancelled lease", err) } return @@ -211,7 +212,7 @@ func TestFinishRunObservationLockTimeoutAndFailureKeepLease(t *testing.T) { if used != nil || code != nil { t.Fatal("failed observation wrote") } - if err = f.writer.CheckExecutionOwnership(t.Context()); err != nil { + if err = f.lease.CheckOwnership(t.Context()); err != nil { t.Fatal("observation cancelled execution lease", err) } }) diff --git a/services/core/internal/execution/owner.go b/services/core/internal/execution/owner.go new file mode 100644 index 00000000..ed34437f --- /dev/null +++ b/services/core/internal/execution/owner.go @@ -0,0 +1,33 @@ +package execution + +import ( + "context" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" +) + +// Ownership is the execution lease as the Worker uses it. *pgunit.Lease implements it. +type Ownership interface { + CheckOwnership(context.Context) error + CancelOperations(context.Context, context.CancelFunc) error + Close(context.Context) error +} + +// Owner is everything bound to one execution lease. Later cutovers add one explicit +// field per domain's execution operations and delete the matching store calls. +type Owner struct { + Lease Ownership + Store *store.Store // the remaining store execution operations, built by store.NewExecution(s, lease) +} + +// leaseCloseTimeout bounds releasing the lease once the Worker owns it. +const leaseCloseTimeout = 5 * time.Second + +// closeLease releases the lease within its own deadline, independent of the +// cancellation of the request or run that ends ownership. +func closeLease(ctx context.Context, lease Ownership) error { + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), leaseCloseTimeout) + defer cancel() + return lease.Close(ctx) +} diff --git a/services/core/internal/execution/owner_test.go b/services/core/internal/execution/owner_test.go new file mode 100644 index 00000000..bdd5ae67 --- /dev/null +++ b/services/core/internal/execution/owner_test.go @@ -0,0 +1,136 @@ +package execution + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" + "github.com/google/uuid" +) + +// closeCountingLease counts Close calls on the lease a Worker owns. It forwards +// to inner; without inner, any call other than Close fails the test. +type closeCountingLease struct { + t *testing.T + inner Ownership + closable atomic.Bool + closes atomic.Int32 +} + +func (l *closeCountingLease) CheckOwnership(ctx context.Context) error { + if l.inner == nil { + l.t.Error("unexpected call to CheckOwnership") + return errors.New("unexpected call to CheckOwnership") + } + return l.inner.CheckOwnership(ctx) +} + +func (l *closeCountingLease) CancelOperations(ctx context.Context, cancel context.CancelFunc) error { + if l.inner == nil { + l.t.Error("unexpected call to CancelOperations") + return errors.New("unexpected call to CancelOperations") + } + return l.inner.CancelOperations(ctx, cancel) +} + +func (l *closeCountingLease) Close(ctx context.Context) error { + l.closes.Add(1) + if !l.closable.Load() { + l.t.Error("lease closed before its owner finished") + } + if _, bounded := ctx.Deadline(); !bounded || ctx.Err() != nil { + l.t.Error("lease closed without a live bounded context", ctx.Err()) + } + if l.inner == nil { + return nil + } + return l.inner.Close(ctx) +} + +func TestStartWorkerFailureClosesLeaseOnce(t *testing.T) { + if _, err := StartWorker(t.Context(), &Dispatcher{}, Owner{}); err == nil { + t.Fatal("worker started without an execution lease") + } + // The failed request's context is already canceled; the close must not inherit it. + canceled, cancel := context.WithCancel(t.Context()) + cancel() + for name, start := range map[string]func(*testing.T, *closeCountingLease) error{ + "negative concurrency": func(t *testing.T, lease *closeCountingLease) error { + _, err := StartWorker(canceled, &Dispatcher{MaxConcurrentExecutions: -1}, Owner{Lease: lease}) + return err + }, + "excess concurrency": func(t *testing.T, lease *closeCountingLease) error { + _, err := StartWorker(canceled, &Dispatcher{MaxConcurrentExecutions: 1025}, Owner{Lease: lease}) + return err + }, + "missing Store": func(t *testing.T, lease *closeCountingLease) error { + _, err := StartWorker(canceled, &Dispatcher{}, Owner{Lease: lease}) + return err + }, + "deployment claim": func(t *testing.T, lease *closeCountingLease) error { + s, owner := resetManagerStore(t) + lease.inner = owner.Lease + id := uuid.NewString() + dispatcher := &Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })} + _, err := StartWorker(canceled, dispatcher, Owner{Lease: lease, Store: owner.Store}) + if ping := owner.Lease.CheckOwnership(t.Context()); !errors.Is(ping, pgunit.ErrLeaseClosed) { + t.Error("failed start kept the database lease", ping) + } + return err + }, + } { + t.Run(name, func(t *testing.T) { + lease := &closeCountingLease{t: t} + lease.closable.Store(true) + if err := start(t, lease); err == nil { + t.Fatal("worker started") + } + if closes := lease.closes.Load(); closes != 1 { + t.Fatal("failed start closed the lease", closes, "times") + } + }) + } +} + +func TestWorkerRunClosesLeaseAfterDrain(t *testing.T) { + s, owner := resetManagerStore(t) + lease := &closeCountingLease{t: t, inner: owner.Lease} + id := uuid.NewString() + dispatcher := &Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })} + worker, err := StartWorker(t.Context(), dispatcher, Owner{Lease: lease, Store: owner.Store}) + if err != nil { + t.Fatal(err) + } + // An external provisioning caller is still in flight when Run exits. + worker.runtimes.active.Add(1) + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { done <- worker.Run(ctx) }() + cancel() + select { + case <-worker.runtimes.ctx.Done(): + case <-time.After(5 * time.Second): + worker.runtimes.active.Done() + t.Fatal("Run did not stop its runtimes") + } + lease.closable.Store(true) + worker.runtimes.active.Done() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatal(err) + } + case <-time.After(5 * time.Second): + t.Fatal("Run did not exit after draining") + } + if closes := lease.closes.Load(); closes != 1 { + t.Fatal("Run closed the lease", closes, "times") + } + if ping := owner.Lease.CheckOwnership(t.Context()); !errors.Is(ping, pgunit.ErrLeaseClosed) { + t.Fatal("Run kept the database lease", ping) + } +} diff --git a/services/core/internal/execution/prepared_dispatch.go b/services/core/internal/execution/prepared_dispatch.go index 6a59799b..b2545dca 100644 --- a/services/core/internal/execution/prepared_dispatch.go +++ b/services/core/internal/execution/prepared_dispatch.go @@ -16,9 +16,10 @@ type EnvironmentRun struct { Turn store.Turn } -// RunEnvironmentInput reserves a Turn on the Session-owned Runtime Executor. -func (d *Dispatcher) RunEnvironmentInput(ctx context.Context, tenantID, sessionID, reservationID string) (run EnvironmentRun, err error) { - if err = d.Store.CheckExecutionOwnership(ctx); err != nil { +// RunEnvironmentInput reserves a Turn on the Session-owned Runtime Executor. It +// checks lease, the lease d.Store was built on, before any Runtime preparation. +func (d *Dispatcher) RunEnvironmentInput(ctx context.Context, lease Ownership, tenantID, sessionID, reservationID string) (run EnvironmentRun, err error) { + if err = lease.CheckOwnership(ctx); err != nil { return run, err } run.Reservation, err = d.Store.ExpireEnvironmentInput(ctx, tenantID, sessionID, reservationID) diff --git a/services/core/internal/execution/runtime_cancellation.go b/services/core/internal/execution/runtime_cancellation.go index dcf62778..7e3dd646 100644 --- a/services/core/internal/execution/runtime_cancellation.go +++ b/services/core/internal/execution/runtime_cancellation.go @@ -14,12 +14,7 @@ func (m *runtimeManager) cancelLifecycles(nodes []*runtimeNode) error { n.lifecycle.cancelOperations() } } - // Managers in pure lifecycle tests have no Store or execution connection. - if m.store == nil { - cancel() - return nil - } - err := m.store.CancelExecutionOperations(m.ctx, cancel) + err := m.lease.CancelOperations(m.ctx, cancel) if err != nil { // A manual reconcile may have no running coordinator to consume failed. // Close admission synchronously; Worker shutdown still owns cancellation. diff --git a/services/core/internal/execution/runtime_compute.go b/services/core/internal/execution/runtime_compute.go index 055d0e2b..42bac479 100644 --- a/services/core/internal/execution/runtime_compute.go +++ b/services/core/internal/execution/runtime_compute.go @@ -94,7 +94,7 @@ func (r *runtimeLifecycle) observeCompute(ctx context.Context, owner store.Runti if owner.SessionDeleted || owner.Expired || owner.State == "cleanup_pending" { return r.cleanupCompute(ctx, p, owner, state) } - if err := r.store.CheckExecutionOwnership(ctx); err != nil { + if err := r.lease.CheckOwnership(ctx); err != nil { return err } switch owner.ComputePhase { diff --git a/services/core/internal/execution/runtime_compute_wake.go b/services/core/internal/execution/runtime_compute_wake.go index 9eae10eb..3692f358 100644 --- a/services/core/internal/execution/runtime_compute_wake.go +++ b/services/core/internal/execution/runtime_compute_wake.go @@ -70,7 +70,7 @@ func (r *runtimeLifecycle) wakeCompute(ctx context.Context, p sandbox.Checkpoint } func (r *runtimeLifecycle) cleanupCompute(ctx context.Context, p sandbox.CheckpointProvider, owner store.RuntimeAllocation, state runtimeCompute) error { - if err := r.store.CheckExecutionOwnership(ctx); err != nil { + if err := r.lease.CheckOwnership(ctx); err != nil { return err } // An uncommitted artifact is found by its persisted attempt, never a directory diff --git a/services/core/internal/execution/runtime_initialization.go b/services/core/internal/execution/runtime_initialization.go index 8db47fce..8d7147a0 100644 --- a/services/core/internal/execution/runtime_initialization.go +++ b/services/core/internal/execution/runtime_initialization.go @@ -130,7 +130,7 @@ func (w *Worker) prepareEnvironment(ctx context.Context, owner store.Environment operations := setupOperations(setup) for index := 0; index < len(cfg.Files)+len(operations); index++ { step, stop := context.WithTimeout(ctx, 2*time.Minute) - err = w.dispatcher.Store.CheckExecutionOwnership(step) + err = w.lease.CheckOwnership(step) if err == nil { var currentPeer = peer currentPeer, err = w.dispatcher.authorizedPeer(step, owner.DeviceID) diff --git a/services/core/internal/execution/runtime_lifecycle.go b/services/core/internal/execution/runtime_lifecycle.go index b09dda19..95643dcd 100644 --- a/services/core/internal/execution/runtime_lifecycle.go +++ b/services/core/internal/execution/runtime_lifecycle.go @@ -41,6 +41,7 @@ type RuntimeProvider struct { type runtimeLifecycle struct { store *store.Store + lease Ownership registry *runtimegateway.Registry config RuntimeProvider nodeID string @@ -55,7 +56,7 @@ type runtimeLifecycle struct { wakeHints chan struct{} } -func newRuntimeManager(s *store.Store, registry *runtimegateway.Registry, config *RuntimeProvider) (*runtimeManager, error) { +func newRuntimeManager(owner Owner, registry *runtimegateway.Registry, config *RuntimeProvider) (*runtimeManager, error) { if config == nil { return nil, nil } @@ -73,7 +74,7 @@ func newRuntimeManager(s *store.Store, registry *runtimegateway.Registry, config } } ctx, stop := context.WithCancel(context.Background()) - return &runtimeManager{store: s, registry: registry, config: copied, setupInstallationID: config.InstallationID, loadDeployment: config.loadDeployment, prepareDeployment: config.prepareDeployment, publishUnconfigured: config.PublishUnconfigured, setupGate: make(chan struct{}, 1), mutationGate: make(chan struct{}, 1), ctx: ctx, cancel: stop, nodes: make(map[string]*runtimeNode), failed: make(chan error, 1), inventory: make(chan struct{}, 1)}, nil + return &runtimeManager{store: owner.Store, lease: owner.Lease, registry: registry, config: copied, setupInstallationID: config.InstallationID, loadDeployment: config.loadDeployment, prepareDeployment: config.prepareDeployment, publishUnconfigured: config.PublishUnconfigured, setupGate: make(chan struct{}, 1), mutationGate: make(chan struct{}, 1), ctx: ctx, cancel: stop, nodes: make(map[string]*runtimeNode), failed: make(chan error, 1), inventory: make(chan struct{}, 1)}, nil } func validatedRuntimeProvider(config *RuntimeProvider, registry *runtimegateway.Registry) (RuntimeProvider, error) { @@ -212,7 +213,7 @@ func (r *runtimeLifecycle) provision(ctx context.Context, tenant, environment, p if owner.Replayed { return owner, nil } - if err := r.store.CheckExecutionOwnership(ctx); err != nil { + if err := r.lease.CheckOwnership(ctx); err != nil { return owner, err } info, err := provider.Create(ctx, sandbox.Bootstrap{ @@ -294,7 +295,7 @@ func (r *runtimeLifecycle) reconcile(ctx context.Context) error { r.recordObservation(ctx, owner, err) stop() if err != nil { - if ownership := r.store.CheckExecutionOwnership(ctx); ownership != nil { + if ownership := r.lease.CheckOwnership(ctx); ownership != nil { return ownership } // Provider errors can include operator configuration. Log safe identity @@ -337,7 +338,7 @@ func (r *runtimeLifecycle) observe(ctx context.Context, owner store.RuntimeAlloc if owner.ProviderKey != r.config.InstallationID { return sandbox.ErrInvalid } - if err := r.store.CheckExecutionOwnership(ctx); err != nil { + if err := r.lease.CheckOwnership(ctx); err != nil { return err } info, err := provider.GetInfo(ctx, runtimeReference(owner)) @@ -366,7 +367,7 @@ func (r *runtimeLifecycle) observe(ctx context.Context, owner store.RuntimeAlloc if err != nil { return err } - if err := r.store.CheckExecutionOwnership(ctx); err != nil { + if err := r.lease.CheckOwnership(ctx); err != nil { return err } if err := provider.Kill(ctx, runtimeReference(owner)); err != nil { diff --git a/services/core/internal/execution/runtime_manager.go b/services/core/internal/execution/runtime_manager.go index 4851bbc9..18f0d874 100644 --- a/services/core/internal/execution/runtime_manager.go +++ b/services/core/internal/execution/runtime_manager.go @@ -17,6 +17,7 @@ var errRuntimeTransition = fmt.Errorf("%w: sandbox configuration is changing", E type runtimeManager struct { store *store.Store + lease Ownership registry *runtimegateway.Registry config RuntimeProvider setupInstallationID string @@ -81,7 +82,7 @@ func (m *runtimeManager) node(id string) (*runtimeNode, error) { if n == nil { ctx, stop := context.WithCancel(m.ctx) n = &runtimeNode{lifecycle: &runtimeLifecycle{ - store: m.store, registry: m.registry, config: m.config, nodeID: id, + store: m.store, lease: m.lease, registry: m.registry, config: m.config, nodeID: id, gate: make(chan struct{}, 1), ctx: ctx, stop: stop, connections: make(map[string]*runtimeConnection), wakeHints: make(chan struct{}, 1), }} diff --git a/services/core/internal/execution/runtime_manager_test.go b/services/core/internal/execution/runtime_manager_test.go index 9fc6b82e..bce4dc07 100644 --- a/services/core/internal/execution/runtime_manager_test.go +++ b/services/core/internal/execution/runtime_manager_test.go @@ -8,10 +8,21 @@ import ( "time" ) +// heldLease is an execution lease that stays held, for lifecycle tests that run +// no execution query: its cancellation fence cancels immediately. +type heldLease struct{} + +func (heldLease) CheckOwnership(context.Context) error { return nil } +func (heldLease) CancelOperations(_ context.Context, cancel context.CancelFunc) error { + cancel() + return nil +} +func (heldLease) Close(context.Context) error { return nil } + func testRuntimeManager(t *testing.T) *runtimeManager { t.Helper() ctx, cancel := context.WithCancel(t.Context()) - m := &runtimeManager{config: RuntimeProvider{ProviderKind: "docker"}, ctx: ctx, cancel: cancel, nodes: make(map[string]*runtimeNode), failed: make(chan error, 1), inventory: make(chan struct{}, 1)} + m := &runtimeManager{lease: heldLease{}, config: RuntimeProvider{ProviderKind: "docker"}, ctx: ctx, cancel: cancel, nodes: make(map[string]*runtimeNode), failed: make(chan error, 1), inventory: make(chan struct{}, 1)} t.Cleanup(func() { m.stop(); m.drain() }) return m } diff --git a/services/core/internal/execution/runtime_pending.go b/services/core/internal/execution/runtime_pending.go index 4f6950e9..ebfb65a6 100644 --- a/services/core/internal/execution/runtime_pending.go +++ b/services/core/internal/execution/runtime_pending.go @@ -28,7 +28,7 @@ func (r *runtimeLifecycle) provisionPending(ctx context.Context) error { _, err := r.provision(operation, environment.TenantID, environment.ID, provider) cancel() if err != nil { - if ownership := r.store.CheckExecutionOwnership(ctx); ownership != nil { + if ownership := r.lease.CheckOwnership(ctx); ownership != nil { return ownership } log.Ctx(ctx).Warn("managed Runtime bootstrap incomplete", "environment_id", environment.ID) diff --git a/services/core/internal/execution/runtime_retirement_failure_test.go b/services/core/internal/execution/runtime_retirement_failure_test.go index dd7418f8..b0ab36f7 100644 --- a/services/core/internal/execution/runtime_retirement_failure_test.go +++ b/services/core/internal/execution/runtime_retirement_failure_test.go @@ -19,9 +19,9 @@ func TestFailedInventoryRetirementClosesAdmissionAndRetainsGate(t *testing.T) { var readOnce sync.Once unblockRead := func() { readOnce.Do(func() { close(releaseRead) }) } defer unblockRead() - writer, pool := delayedReadWriter(t, &armed, reading, releaseRead) + owner, pool := delayedReadWriter(t, &armed, reading, releaseRead) m := testRuntimeManager(t) - m.store = writer + m.store, m.lease = owner.Store, owner.Lease m.loadDeployment = func(context.Context) (*RuntimeProvider, error) { return nil, nil } m.mutationGate = make(chan struct{}, 1) // This fixture models an already loaded node deployment; its provider is @@ -49,7 +49,7 @@ func TestFailedInventoryRetirementClosesAdmissionAndRetainsGate(t *testing.T) { defer cancel() queryDone = make(chan error, 1) armed.Store(true) - go func() { queryDone <- writer.CheckExecutionOwnership(queryCtx) }() + go func() { queryDone <- owner.Lease.CheckOwnership(queryCtx) }() select { case <-reading: case <-time.After(2 * time.Second): diff --git a/services/core/internal/execution/sandbox_deployment_drain_test.go b/services/core/internal/execution/sandbox_deployment_drain_test.go index 8c7a9477..436c61ec 100644 --- a/services/core/internal/execution/sandbox_deployment_drain_test.go +++ b/services/core/internal/execution/sandbox_deployment_drain_test.go @@ -11,6 +11,7 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/node" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" @@ -45,7 +46,7 @@ func (c *delayedLeaseRead) Read(p []byte) (int, error) { // delayedReadWriter owns an isolated database, so each case has its own // advisory-lock namespace. The delayed driver read lets a cancellation fence hit // its own deadline without shortening production timeouts. -func delayedReadWriter(t *testing.T, armed *atomic.Bool, reading chan struct{}, release <-chan struct{}) (*store.Store, *pgxpool.Pool) { +func delayedReadWriter(t *testing.T, armed *atomic.Bool, reading chan struct{}, release <-chan struct{}) (Owner, *pgxpool.Pool) { t.Helper() pool := pgtest.OpenIsolated(t, func(cfg *pgxpool.Config) { dial := cfg.ConnConfig.DialFunc @@ -57,18 +58,18 @@ func delayedReadWriter(t *testing.T, armed *atomic.Bool, reading chan struct{}, return &delayedLeaseRead{Conn: c, armed: armed, reading: reading, release: release}, nil } }) - writer, err := store.NewExecution(t.Context(), store.New(pool)) + lease, err := pgunit.AcquireLease(t.Context(), pool) if err != nil { t.Fatal(err) } t.Cleanup(func() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if err := writer.CloseExecution(ctx); err != nil { + if err := lease.Close(ctx); err != nil { t.Error(err) } }) - return writer, pool + return Owner{Lease: lease, Store: store.NewExecution(store.New(pool), lease)}, pool } func TestSandboxDeploymentDrainPreservesLeaseInFlightRead(t *testing.T) { @@ -85,12 +86,12 @@ func testLifecycleCancellationPreservesLease(t *testing.T, mode string) { var releaseOnce sync.Once unblock := func() { releaseOnce.Do(func() { close(release) }) } defer unblock() - writer, _ := delayedReadWriter(t, &armed, reading, release) + owner, _ := delayedReadWriter(t, &armed, reading, release) hub := node.NewHub(node.HubOptions{}) defer hub.Close() id := uuid.NewString() configuration := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", 1)} - m, err := newRuntimeManager(writer, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return configuration, nil })) + m, err := newRuntimeManager(owner, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return configuration, nil })) if err != nil { t.Fatal(err) } @@ -124,11 +125,11 @@ func testLifecycleCancellationPreservesLease(t *testing.T, mode string) { return } defer finish() - queryDone <- writer.CheckExecutionOwnership(ctx) + queryDone <- owner.Lease.CheckOwnership(ctx) <-ctx.Done() // Provider settlement remains outside the cancellation fence. A fresh owner // read must proceed even before this tracked lifecycle operation returns. - leaseFree <- writer.CheckExecutionOwnership(t.Context()) + leaseFree <- owner.Lease.CheckOwnership(t.Context()) }() select { case <-reading: @@ -181,7 +182,7 @@ func testLifecycleCancellationPreservesLease(t *testing.T, mode string) { if delayed != nil { delayed.unblock() } - if err := writer.CheckExecutionOwnership(t.Context()); err != nil { + if err := owner.Lease.CheckOwnership(t.Context()); err != nil { t.Fatalf("deployment drain destroyed the owner connection (in-flight query: %v): %v", queryErr, err) } if queryErr != nil { @@ -223,12 +224,12 @@ func (c *delayedCancellationContext) unblock() { } func TestSandboxDeploymentDrainFailureCannotReactivate(t *testing.T) { - _, writer := resetManagerStore(t) + _, owner := resetManagerStore(t) hub := node.NewHub(node.HubOptions{}) defer hub.Close() id := uuid.NewString() configuration := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", 1)} - m, err := newRuntimeManager(writer, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return configuration, nil })) + m, err := newRuntimeManager(owner, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return configuration, nil })) if err != nil { t.Fatal(err) } @@ -240,7 +241,7 @@ func TestSandboxDeploymentDrainFailureCannotReactivate(t *testing.T) { if err != nil { t.Fatal(err) } - if err := writer.CloseExecution(t.Context()); err != nil { + if err := owner.Lease.Close(t.Context()); err != nil { t.Fatal(err) } first := m.pauseDeployment(t.Context()) diff --git a/services/core/internal/execution/sandbox_deployment_setup_test.go b/services/core/internal/execution/sandbox_deployment_setup_test.go index 084087c8..50624769 100644 --- a/services/core/internal/execution/sandbox_deployment_setup_test.go +++ b/services/core/internal/execution/sandbox_deployment_setup_test.go @@ -25,7 +25,7 @@ func TestDeferredSandboxDeploymentLoadsOnceBeforeNodeCreation(t *testing.T) { var selected atomic.Bool var loads atomic.Int32 configuration := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", 1)} - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { loads.Add(1) if !selected.Load() { return nil, nil @@ -70,7 +70,7 @@ func TestDeferredSandboxDeploymentLoadsOnceBeforeNodeCreation(t *testing.T) { func TestDeferredSandboxDeploymentShutdownCancelsLoad(t *testing.T) { entered := make(chan struct{}) - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(uuid.NewString(), func(ctx context.Context) (*RuntimeProvider, error) { + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(uuid.NewString(), func(ctx context.Context) (*RuntimeProvider, error) { close(entered) <-ctx.Done() return nil, ctx.Err() @@ -95,7 +95,7 @@ func TestDeferredSandboxProviderFailureKeepsRecoveryAvailable(t *testing.T) { available := false loadErr := ErrExecutionUnavailable configuration := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", 1)} - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { if !available { return nil, loadErr } @@ -127,7 +127,7 @@ func TestRejectedSandboxCandidatePreservesActiveGeneration(t *testing.T) { id := uuid.NewString() config := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", 1)} rejected := errors.New("candidate provider unavailable") - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil }, + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil }, func(context.Context, store.SandboxSetup) (PreparedRuntimeDeployment, error) { return PreparedRuntimeDeployment{}, rejected })) @@ -159,7 +159,7 @@ func TestRejectedSandboxCandidatePreservesActiveGeneration(t *testing.T) { func TestSandboxCandidateValidationDoesNotHoldManagerLock(t *testing.T) { id := uuid.NewString() entered, release := make(chan struct{}), make(chan struct{}) - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil }, + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil }, func(context.Context, store.SandboxSetup) (PreparedRuntimeDeployment, error) { close(entered) <-release @@ -192,7 +192,7 @@ func TestCommittedSandboxCandidatePublishesAfterShutdown(t *testing.T) { hub := node.NewHub(node.HubOptions{}) defer hub.Close() id := uuid.NewString() - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_deployment_switch_test.go b/services/core/internal/execution/sandbox_deployment_switch_test.go index ae14dac0..6ed04d45 100644 --- a/services/core/internal/execution/sandbox_deployment_switch_test.go +++ b/services/core/internal/execution/sandbox_deployment_switch_test.go @@ -18,7 +18,7 @@ func TestSandboxManagerSwitchDrainsBeforeDirectActivation(t *testing.T) { defer hub.Close() id := uuid.NewString() config := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", 1)} - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) if err != nil { t.Fatal(err) } @@ -79,7 +79,7 @@ func TestSandboxManagerSwitchDrainsBeforeDirectActivation(t *testing.T) { func TestSandboxManagerFailedActivationStaysPaused(t *testing.T) { id := uuid.NewString() - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, errors.New("provider unavailable") })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, errors.New("provider unavailable") })) if err != nil { t.Fatal(err) } @@ -100,7 +100,7 @@ func TestSandboxManagerCancelledSwitchCannotResumeBeforeDrain(t *testing.T) { defer hub.Close() id := uuid.NewString() config := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", 1)} - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) if err != nil { t.Fatal(err) } @@ -154,7 +154,7 @@ func TestSandboxActivationCannotBypassOutstandingDrain(t *testing.T) { hub := node.NewHub(node.HubOptions{}) defer hub.Close() id := uuid.NewString() - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil }, func(_ context.Context, setup store.SandboxSetup) (PreparedRuntimeDeployment, error) { return PreparedRuntimeDeployment{Config: &RuntimeProvider{InstallationID: setup.InstallationID, ProviderKind: setup.Provider, Mode: setup.Mode, CoreURL: "https://core.example/api/v1", BackendFingerprint: setup.BackendFingerprint, Provider: hub.Proxy(uuid.NewString(), "docker", 1)}}, nil diff --git a/services/core/internal/execution/sandbox_generations_test.go b/services/core/internal/execution/sandbox_generations_test.go index 7bdc2c78..feb60b4c 100644 --- a/services/core/internal/execution/sandbox_generations_test.go +++ b/services/core/internal/execution/sandbox_generations_test.go @@ -15,7 +15,8 @@ import ( ) func TestE2BReplacementVerifiesTwiceAndNeverPublishesFailedCommit(t *testing.T) { - s, writer := resetManagerStore(t) + s, owner := resetManagerStore(t) + writer := owner.Store id := uuid.NewString() if err := writer.ClaimWebSandboxDeployment(t.Context(), id); err != nil { t.Fatal(err) @@ -49,7 +50,7 @@ func TestE2BReplacementVerifiesTwiceAndNeverPublishesFailedCommit(t *testing.T) }, FenceCredential: func(context.Context) (func(), error) { fenced++; return func() { released++ }, nil }, Publish: func(*RuntimeProvider) { published++ }}, nil }) - m, err := newRuntimeManager(writer, runtimegateway.NewRegistry(), config) + m, err := newRuntimeManager(owner, runtimegateway.NewRegistry(), config) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_provider_contract_test.go b/services/core/internal/execution/sandbox_provider_contract_test.go index a73cd482..b84b05c9 100644 --- a/services/core/internal/execution/sandbox_provider_contract_test.go +++ b/services/core/internal/execution/sandbox_provider_contract_test.go @@ -18,7 +18,7 @@ func TestSandboxProviderRegistrationDoesNotRequireAnExecutionVendorBranch(t *tes id := uuid.NewString() config := &RuntimeProvider{InstallationID: id, ProviderKind: "contract-fixture", Mode: mode, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: &lifecycleOnlySandbox{}} - m, err := newRuntimeManager(nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_reset_test.go b/services/core/internal/execution/sandbox_reset_test.go index 2fdcfc88..1ebdca6b 100644 --- a/services/core/internal/execution/sandbox_reset_test.go +++ b/services/core/internal/execution/sandbox_reset_test.go @@ -11,6 +11,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/adminaudit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/node" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" @@ -19,13 +20,14 @@ import ( ) // The execution lease is database-scoped, so these manager tests own a database. -// They receive the pooled Store and the execution writer built on it. -func resetManagerStore(t *testing.T) (*store.Store, *store.Store) { +// They receive the pooled Store and the Owner of its execution lease, which the +// test closes when it ends. +func resetManagerStore(t *testing.T) (*store.Store, Owner) { t.Helper() return resetManagerStoreConfig(t, nil) } -func resetManagerStoreConfig(t *testing.T, configure func(*pgxpool.Config)) (*store.Store, *store.Store) { +func resetManagerStoreConfig(t *testing.T, configure func(*pgxpool.Config)) (*store.Store, Owner) { t.Helper() pool := pgtest.OpenIsolated(t, configure) cipher, err := credentialcrypto.New(bytes.Repeat([]byte{8}, 32)) @@ -33,16 +35,24 @@ func resetManagerStoreConfig(t *testing.T, configure func(*pgxpool.Config)) (*st t.Fatal(err) } s := store.NewWithCredentialCipher(pool, cipher) - writer, err := store.NewExecution(t.Context(), s) + return s, testOwner(t, pool, s) +} + +// testOwner acquires the execution lease on pool and builds s's execution +// writer on it, as cmd/server does. The lease closes when the test ends. +func testOwner(t *testing.T, pool *pgxpool.Pool, s *store.Store) Owner { + t.Helper() + lease, err := pgunit.AcquireLease(t.Context(), pool) if err != nil { t.Fatal(err) } - t.Cleanup(func() { _ = writer.CloseExecution(context.Background()) }) - return s, writer + t.Cleanup(func() { _ = lease.Close(context.Background()) }) + return Owner{Lease: lease, Store: store.NewExecution(s, lease)} } func TestSandboxResetPageTimeoutRecoversCommittedOwner(t *testing.T) { - s, w := resetManagerStore(t) + s, owner := resetManagerStore(t) + w := owner.Store id := uuid.NewString() if err := w.ClaimWebSandboxDeployment(t.Context(), id); err != nil { t.Fatal(err) @@ -67,7 +77,7 @@ func TestSandboxResetPageTimeoutRecoversCommittedOwner(t *testing.T) { } return &RuntimeProvider{InstallationID: id, ProviderKind: setup.Provider, Mode: setup.Mode, Generation: setup.Generation, CoreURL: "https://core.example/api/v1", BackendFingerprint: setup.BackendFingerprint, Provider: hub.Proxy(uuid.NewString(), "docker", 1)}, nil }) - m, err := newRuntimeManager(w, runtimegateway.NewRegistry(), config) + m, err := newRuntimeManager(owner, runtimegateway.NewRegistry(), config) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_snapshot_budget_test.go b/services/core/internal/execution/sandbox_snapshot_budget_test.go index 596a2913..164f5f58 100644 --- a/services/core/internal/execution/sandbox_snapshot_budget_test.go +++ b/services/core/internal/execution/sandbox_snapshot_budget_test.go @@ -60,10 +60,11 @@ func (d *snapshotBudget) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, data pg func TestSandboxResetSnapshotFitsPageBudget(t *testing.T) { budget := &snapshotBudget{t: t} - s, w := resetManagerStoreConfig(t, func(cfg *pgxpool.Config) { + s, owner := resetManagerStoreConfig(t, func(cfg *pgxpool.Config) { cfg.ConnConfig.RuntimeParams["jit"] = "on" cfg.ConnConfig.Tracer = budget }) + w := owner.Store id := uuid.NewString() if err := w.ClaimWebSandboxDeployment(t.Context(), id); err != nil { t.Fatal(err) @@ -83,7 +84,7 @@ func TestSandboxResetSnapshotFitsPageBudget(t *testing.T) { } return &RuntimeProvider{InstallationID: id, ProviderKind: setup.Provider, Generation: setup.Generation, Mode: setup.Mode, CoreURL: "https://core.example/api/v1", BackendFingerprint: setup.BackendFingerprint, Provider: hub.Proxy(uuid.NewString(), "docker", 1)}, nil }) - m, err := newRuntimeManager(w, runtimegateway.NewRegistry(), configuration) + m, err := newRuntimeManager(owner, runtimegateway.NewRegistry(), configuration) if err != nil { t.Fatal(err) } @@ -97,7 +98,7 @@ func TestSandboxResetSnapshotFitsPageBudget(t *testing.T) { budget.armed.Store(true) started := time.Now() err = m.resetStep(t.Context()) - ping := w.CheckExecutionOwnership(t.Context()) + ping := owner.Lease.CheckOwnership(t.Context()) t.Logf("reset_elapsed=%s reset_error=%v lease_ping=%v", time.Since(started), err, ping) if err != nil || ping != nil { t.Fatal("bounded snapshot lost execution ownership", err, ping) diff --git a/services/core/internal/execution/worker.go b/services/core/internal/execution/worker.go index fd4c5efc..0d8989db 100644 --- a/services/core/internal/execution/worker.go +++ b/services/core/internal/execution/worker.go @@ -20,6 +20,7 @@ type Worker struct { metrics workerMetricsState dispatcher *Dispatcher admission *store.Store + lease Ownership directoryReads chan directoryReadRequest fileWrites chan fileWriteRequest scheduleWake chan struct{} @@ -29,23 +30,42 @@ type Worker struct { enrolledConnections map[string]*runtimeConnection } -func StartWorker(ctx context.Context, dispatcher *Dispatcher) (*Worker, error) { +// StartWorker takes over owner.Lease from the moment it is called: a failed +// start closes the lease before returning, and a started Worker closes it after +// Run drains. +func StartWorker(ctx context.Context, dispatcher *Dispatcher, owner Owner) (_ *Worker, err error) { + if owner.Lease == nil { + return nil, errors.New("execution worker requires an execution lease") + } + defer func() { + if err == nil { + return + } + if closeErr := closeLease(ctx, owner.Lease); closeErr != nil { + err = errors.Join(err, closeErr) + } + }() if dispatcher.MaxConcurrentExecutions < 0 || dispatcher.MaxConcurrentExecutions > 1024 { return nil, errors.New("execution concurrency must be between 1 and 1024, or zero for the default") } - writer, err := store.NewExecution(ctx, dispatcher.Store) - if err != nil { - return nil, err + if owner.Store == nil { + return nil, errors.New("execution worker requires the execution Store") } owned := *dispatcher - owned.Store = writer + owned.Store = owner.Store owned.notifications = &executionNotifications{} - worker := &Worker{concurrency: dispatcher.MaxConcurrentExecutions, dispatcher: &owned, admission: dispatcher.Store, directoryReads: make(chan directoryReadRequest), fileWrites: make(chan fileWriteRequest), stopped: make(chan struct{}), scheduleWake: make(chan struct{}, 1), enrolledConnections: make(map[string]*runtimeConnection)} - worker.runtimes, err = newRuntimeManager(owned.Store, owned.Registry, owned.ManagedRuntimes) + worker := &Worker{concurrency: dispatcher.MaxConcurrentExecutions, dispatcher: &owned, admission: dispatcher.Store, lease: owner.Lease, directoryReads: make(chan directoryReadRequest), fileWrites: make(chan fileWriteRequest), stopped: make(chan struct{}), scheduleWake: make(chan struct{}, 1), enrolledConnections: make(map[string]*runtimeConnection)} + worker.runtimes, err = newRuntimeManager(owner, owned.Registry, owned.ManagedRuntimes) if err != nil { - _ = writer.CloseExecution(context.Background()) return nil, err } + if worker.runtimes != nil { + defer func() { + if err != nil { + worker.runtimes.stop() + } + }() + } var deployment *store.RuntimeDeployment if worker.runtimes != nil && worker.runtimes.loadDeployment == nil { config := worker.runtimes.config @@ -60,24 +80,12 @@ func StartWorker(ctx context.Context, dispatcher *Dispatcher) (*Worker, error) { err = owned.Store.ConfigureRuntimeDeployment(ctx, deployment) } if err != nil { - if worker.runtimes != nil { - worker.runtimes.stop() - } - _ = writer.CloseExecution(context.Background()) return nil, err } - if err := owned.Store.ReconcileEnvironmentConnections(ctx); err != nil { - if worker.runtimes != nil { - worker.runtimes.stop() - } - _ = writer.CloseExecution(context.Background()) + if err = owned.Store.ReconcileEnvironmentConnections(ctx); err != nil { return nil, err } - if err := worker.reconcile(ctx); err != nil { - if worker.runtimes != nil { - worker.runtimes.stop() - } - _ = writer.CloseExecution(context.Background()) + if err = worker.reconcile(ctx); err != nil { return nil, err } worker.observeOwnership(nil) @@ -86,7 +94,7 @@ func StartWorker(ctx context.Context, dispatcher *Dispatcher) (*Worker, error) { // CheckOwnership checks the same database lease used for execution writes. func (w *Worker) CheckOwnership(ctx context.Context) error { - err := w.dispatcher.Store.CheckExecutionOwnership(ctx) + err := w.lease.CheckOwnership(ctx) w.observeOwnership(err) return err } @@ -148,9 +156,7 @@ func (w *Worker) Run(ctx context.Context) (runErr error) { // Drain an external provisioning caller before releasing the writer lease. w.runtimes.drain() } - closeCtx, stop := context.WithTimeout(context.Background(), 5*time.Second) - defer stop() - w.observeWorkerClosed(w.dispatcher.Store.CloseExecution(closeCtx)) + w.observeWorkerClosed(closeLease(ctx, w.lease)) }() active := make(map[string]bool) w.observeSlots(len(active)) diff --git a/services/core/internal/execution/worker_concurrency_test.go b/services/core/internal/execution/worker_concurrency_test.go index 87b8f7e6..6307f095 100644 --- a/services/core/internal/execution/worker_concurrency_test.go +++ b/services/core/internal/execution/worker_concurrency_test.go @@ -11,10 +11,4 @@ func TestWorkerConfiguredConcurrencyReportsActualCapacity(t *testing.T) { t.Fatal("metrics do not reflect execution concurrency", snapshot) } } - for _, invalid := range []int{-1, 1025} { - // Reject before accessing the Store or acquiring its execution lease. - if _, err := StartWorker(t.Context(), &Dispatcher{MaxConcurrentExecutions: invalid}); err == nil { - t.Fatal("invalid concurrency accepted", invalid) - } - } } diff --git a/services/core/internal/execution/worker_metrics_test.go b/services/core/internal/execution/worker_metrics_test.go index 93b14076..c8dcb374 100644 --- a/services/core/internal/execution/worker_metrics_test.go +++ b/services/core/internal/execution/worker_metrics_test.go @@ -8,7 +8,7 @@ import ( "testing" "time" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" ) func TestWorkerMetricsUnknownAndDetached(t *testing.T) { @@ -39,11 +39,19 @@ func TestWorkerMetricsUnknownAndDetached(t *testing.T) { } } +// lostLease is an execution lease that no longer owns the database. +type lostLease struct{} + +func (lostLease) CheckOwnership(context.Context) error { return pgunit.ErrLeaseClosed } +func (lostLease) CancelOperations(context.Context, context.CancelFunc) error { + return pgunit.ErrLeaseClosed +} +func (lostLease) Close(context.Context) error { return nil } + func TestWorkerMetricsFailuresAndClosure(t *testing.T) { - // A pooled Store has no execution lease, so its ownership check fails. - worker := &Worker{dispatcher: &Dispatcher{Store: &store.Store{}}} + worker := &Worker{lease: lostLease{}} worker.observeOwnership(nil) - if err := worker.CheckOwnership(t.Context()); !errors.Is(err, store.ErrExecutionAuthority) { + if err := worker.CheckOwnership(t.Context()); !errors.Is(err, pgunit.ErrLeaseClosed) { t.Fatalf("ownership error changed: %v", err) } if worker.MetricsSnapshot().ExecutionOwner != nil { diff --git a/services/core/internal/execution/worker_schedule.go b/services/core/internal/execution/worker_schedule.go index 19b19509..70b103d9 100644 --- a/services/core/internal/execution/worker_schedule.go +++ b/services/core/internal/execution/worker_schedule.go @@ -89,7 +89,7 @@ func (s *workerSchedule) selectWork(ctx context.Context, w *Worker, devices []st } func (w *Worker) runEnvironmentInput(ctx context.Context, item scheduledWork) error { - run, err := w.dispatcher.RunEnvironmentInput(ctx, item.TenantID, item.SessionID, item.reservationID) + run, err := w.dispatcher.RunEnvironmentInput(ctx, w.lease, item.TenantID, item.SessionID, item.reservationID) if err == nil { return nil } diff --git a/services/core/internal/sandbox/providers/configuration_flow_test.go b/services/core/internal/sandbox/providers/configuration_flow_test.go index 4031d5c1..d0825d2a 100644 --- a/services/core/internal/sandbox/providers/configuration_flow_test.go +++ b/services/core/internal/sandbox/providers/configuration_flow_test.go @@ -6,6 +6,7 @@ import ( "encoding/json" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/api" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/providercontract" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" @@ -85,11 +86,12 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { adapter.Configuration = regionalCodec{} providers.RegisterFixture(t, kind, adapter) s := store.New(pool) - w, err := store.NewExecution(t.Context(), s) + lease, err := pgunit.AcquireLease(t.Context(), pool) if err != nil { t.Fatal(err) } - defer w.CloseExecution(context.Background()) + defer lease.Close(context.Background()) + w := store.NewExecution(s, lease) installation := uuid.NewString() if err = w.ClaimWebSandboxDeployment(t.Context(), installation); err != nil { t.Fatal(err) diff --git a/services/core/internal/store/admin_session_archive_worker_http_test.go b/services/core/internal/store/admin_session_archive_worker_http_test.go index 240404c0..9ebbebff 100644 --- a/services/core/internal/store/admin_session_archive_worker_http_test.go +++ b/services/core/internal/store/admin_session_archive_worker_http_test.go @@ -32,7 +32,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} installation := uuid.NewString() provider := &lifecycleProvider{resources: map[string]sandbox.Info{}} providerConfig := func(setup store.SandboxSetup) *execution.RuntimeProvider { @@ -47,10 +47,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { }, func(_ context.Context, setup store.SandboxSetup) (execution.PreparedRuntimeDeployment, error) { return execution.PreparedRuntimeDeployment{Config: providerConfig(setup)}, nil }) - worker, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) var once sync.Once stop := func() { once.Do(func() { @@ -94,7 +91,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, s, nil, "codex", storeKeys(s), workerExecution(worker), withCoreKeys(admin)) + handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), workerExecution(worker), withCoreKeys(admin)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/agent_execution_defaults_http_test.go b/services/core/internal/store/agent_execution_defaults_http_test.go index 089dc8b7..d7084c4f 100644 --- a/services/core/internal/store/agent_execution_defaults_http_test.go +++ b/services/core/internal/store/agent_execution_defaults_http_test.go @@ -18,12 +18,12 @@ import ( func TestAgentExecutionDefaultsPublicSnapshotAndPrecedence(t *testing.T) { _, pool := store.NewTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{19}, 32)) - st := store.NewWithCredentialCipher(pool, cipher) + st, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, st), managedSandboxes(t, st), withHarnesses([]string{"codex", "claude_sdk", "mcode"}), modelProviderDefaults(st, func(context.Context, string) (*store.DeploymentModelProviderSnapshot, error) { + handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, st), withHarnesses([]string{"codex", "claude_sdk", "mcode"}), modelProviderDefaults(st, func(context.Context, string) (*store.DeploymentModelProviderSnapshot, error) { defaultsCalls++ copy := *deployment return &store.DeploymentModelProviderSnapshot{Model: "fixture", Provider: ©, Revision: uuid.New()}, nil diff --git a/services/core/internal/store/agents_delete_public_test.go b/services/core/internal/store/agents_delete_public_test.go index dbfc1cce..2ccc36e9 100644 --- a/services/core/internal/store/agents_delete_public_test.go +++ b/services/core/internal/store/agents_delete_public_test.go @@ -18,20 +18,20 @@ func TestAgentDeletionOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, pool := store.NewTestStore(t) + s, db := newTestStoreDB(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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() - recoveredStore := store.New(pool) - h, err = publicHandler(t, recoveredStore, auth, "codex", storeExecution(t, recoveredStore)) + recoveredStore, recoveredDB := store.New(db.pool), fixtureDB{pool: db.pool} + h, err = publicHandler(t, recoveredStore, recoveredDB, auth, "codex", storeExecution(t, recoveredStore)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/agents_update_public_test.go b/services/core/internal/store/agents_update_public_test.go index 42b79090..2930a267 100644 --- a/services/core/internal/store/agents_update_public_test.go +++ b/services/core/internal/store/agents_update_public_test.go @@ -18,20 +18,20 @@ func TestAgentUpdateOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, pool := store.NewTestStore(t) + s, db := newTestStoreDB(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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() - recoveredStore := store.New(pool) - h, err = publicHandler(t, recoveredStore, auth, "codex", storeExecution(t, recoveredStore)) + recoveredStore, recoveredDB := store.New(db.pool), fixtureDB{pool: db.pool} + h, err = publicHandler(t, recoveredStore, recoveredDB, auth, "codex", storeExecution(t, recoveredStore)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/archive_cancellation_test.go b/services/core/internal/store/archive_cancellation_test.go index f38df21a..9cfe0d18 100644 --- a/services/core/internal/store/archive_cancellation_test.go +++ b/services/core/internal/store/archive_cancellation_test.go @@ -31,13 +31,11 @@ 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, pool := store.NewManagedTestStore(t) - writer, err := store.NewExecution(t.Context(), s) - if err != nil { - t.Fatal(err) - } + s, db := newManagedTestStoreDB(t) + leased := executionOwner(t, db, s) + writer := leased.Store t.Cleanup(func() { - if err := writer.CloseExecution(context.Background()); err != nil { + if err := leased.Lease.Close(context.Background()); err != nil { t.Error(err) } }) @@ -96,7 +94,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { t.Fatal(err) } t.Cleanup(func() { conn.Close() }) - h := &dispatchHarness{t: t, s: s, tenant: project.TenantID, session: session, conn: conn, registry: registry, d: &execution.Dispatcher{Store: writer, Registry: registry}} + h := &dispatchHarness{t: t, s: s, db: db, lease: leased.Lease, tenant: project.TenantID, session: session, conn: conn, registry: registry, d: &execution.Dispatcher{Store: writer, Registry: registry}} capabilities := workerEnvironmentCapabilities() capabilities.FunctionTools = proto.CapabilitySupported h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: capabilities}}}) @@ -210,18 +208,18 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { t.Fatal("missing cancel delivery identity") } if scenario == "rotated" { - if _, err := pool.Exec(t.Context(), "UPDATE devices SET credential_hash=$2 WHERE id=$1", owner.DeviceID, runtimedevice.HashCredential(uuid.NewString())); err != nil { + if _, err := db.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 := pool.Exec(t.Context(), "UPDATE turns SET cancel_requested_at=clock_timestamp()-interval '21 seconds' WHERE id=$1", input.TurnID); err != nil { + 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 { t.Fatal(err) } } var unlockCommit func() if scenario == "ack_commit_blocked" { - tx, err := pool.Begin(t.Context()) + tx, err := db.pool.Begin(t.Context()) if err != nil { t.Fatal(err) } @@ -261,7 +259,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 := 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 := 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 { t.Fatal(err) } if blocked { @@ -294,7 +292,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { } var receipts int - if err := 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 := db.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 @@ -311,7 +309,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { t.Fatal(allocation, err) } var revoked bool - if err := pool.QueryRow(t.Context(), "SELECT revoked_at IS NOT NULL FROM devices WHERE id=$1", owner.DeviceID).Scan(&revoked); err != nil || !revoked { + 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 { t.Fatal(revoked, err) } }) diff --git a/services/core/internal/store/claude_execution_test.go b/services/core/internal/store/claude_execution_test.go index 0ab01f37..7353e0a4 100644 --- a/services/core/internal/store/claude_execution_test.go +++ b/services/core/internal/store/claude_execution_test.go @@ -11,7 +11,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) @@ -55,10 +54,7 @@ func TestClaudeWorkerSelectsStoredEngineAndRestrictiveCapabilities(t *testing.T) input := h.message("start", "Look up ticket") ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -126,10 +122,7 @@ func TestClaudeDispatcherRejectsUnsupportedConfigurationBeforeClaim(t *testing.T func TestClaudeInvalidImageResultRejectsWholeBatchBeforePersistence(t *testing.T) { h := newDispatchHarness(t) claudeSession(t, h, functionConfiguration, false) - worker, err := execution.StartWorker(t.Context(), h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), h.db, h.d) defer func() { ctx, cancel := context.WithCancel(context.Background()); cancel(); _ = worker.Run(ctx) }() input := h.message("start", "Run") if _, err := h.s.TransitionTurn(t.Context(), h.tenant, h.session.ID, input.TurnID, store.TurnTransition{ExpectedStatus: store.TurnQueued, Status: store.TurnInProgress}); err != nil { diff --git a/services/core/internal/store/claude_mcp_test.go b/services/core/internal/store/claude_mcp_test.go index 86b164ce..287ed009 100644 --- a/services/core/internal/store/claude_mcp_test.go +++ b/services/core/internal/store/claude_mcp_test.go @@ -10,7 +10,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) @@ -44,10 +43,7 @@ func TestClaudeMCPWaitsForCapableRuntime(t *testing.T) { } ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { diff --git a/services/core/internal/store/configuration_validation_public_test.go b/services/core/internal/store/configuration_validation_public_test.go index d91df934..c7e242ad 100644 --- a/services/core/internal/store/configuration_validation_public_test.go +++ b/services/core/internal/store/configuration_validation_public_test.go @@ -24,13 +24,13 @@ func TestAgentConfigurationValidationRejectsWithoutWritesPostgres(t *testing.T) if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/creation_stream_settlement_public_test.go b/services/core/internal/store/creation_stream_settlement_public_test.go index c77c7afb..430def91 100644 --- a/services/core/internal/store/creation_stream_settlement_public_test.go +++ b/services/core/internal/store/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, pool := store.NewModelTestStore(t) + s, db := newModelTestStoreDB(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, auth, "codex", storeExecution(t, s), executorURL("https://offline-executor.example")) + handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://offline-executor.example")) if err != nil { t.Fatal(err) } @@ -127,11 +127,7 @@ func TestCreationStreamPublicLifetimes(t *testing.T) { handler.ServeHTTP(w, r) })) defer server.Close() - writer, err := store.NewExecution(t.Context(), s) - if err != nil { - t.Fatal(err) - } - defer func() { _ = writer.CloseExecution(context.Background()) }() + writer := executionOwner(t, db, s).Store connect := func(environment string) { t.Helper() generation := uuid.NewString() @@ -207,7 +203,7 @@ func TestCreationStreamPublicLifetimes(t *testing.T) { } fresh.open(t) var reservation string - if err := pool.QueryRow(t.Context(), "SELECT id FROM environment_input_reservations WHERE session_id=$1 AND is_initial", session).Scan(&reservation); err != nil { + 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 { t.Fatal(err) } cursor, err := s.SessionEventCursor(t.Context(), tenant, session) diff --git a/services/core/internal/store/credential_matrix_http_test.go b/services/core/internal/store/credential_matrix_http_test.go index 0b6a84c2..89f76fd7 100644 --- a/services/core/internal/store/credential_matrix_http_test.go +++ b/services/core/internal/store/credential_matrix_http_test.go @@ -24,7 +24,7 @@ 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, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) s.SetPublicURL("https://core.example") ctx := t.Context() coreKey := uuid.NewString() @@ -32,7 +32,7 @@ func TestCredentialNamespaceMatrix(t *testing.T) { if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, s, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, s), withCoreKeys(admin)) + handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, s), withCoreKeys(admin)) if err != nil { t.Fatal(err) } @@ -94,10 +94,7 @@ func TestCredentialNamespaceMatrix(t *testing.T) { }, func(_ context.Context, setup store.SandboxSetup) (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, err := execution.StartWorker(ctx, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: runtimes}) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: runtimes}) var stop sync.Once t.Cleanup(func() { stop.Do(func() { diff --git a/services/core/internal/store/deployment_model_providers_http_test.go b/services/core/internal/store/deployment_model_providers_http_test.go index 0332f795..182325d9 100644 --- a/services/core/internal/store/deployment_model_providers_http_test.go +++ b/services/core/internal/store/deployment_model_providers_http_test.go @@ -29,14 +29,14 @@ func TestDeploymentModelProvidersHTTP(t *testing.T) { t.Fatal(err) } cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{53}, 32)) - st := store.NewWithCredentialCipher(pool, cipher) + st, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, st), managedSandboxes(t, st), withCoreKeys(admin), withHarnesses([]string{"codex", "mcode"})) + handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, st), withCoreKeys(admin), withHarnesses([]string{"codex", "mcode"})) if err != nil { t.Fatal(err) } @@ -240,7 +240,7 @@ func TestLegacySessionWithoutProviderCannotStartWork(t *testing.T) { if err != nil { t.Fatal(err) } - worker, stop := startEnvironmentExpiryWorker(t, h.d) + worker, stop := startEnvironmentExpiryWorker(t, h.db, h.d) defer stop() _, pool := store.NewTestStore(t) reservations := func() int { @@ -283,13 +283,13 @@ 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, pool := store.NewModelTestStore(t) - if _, err := pool.Exec(t.Context(), "DELETE FROM deployment_model_providers"); err != nil { + st, db := newModelTestStoreDB(t) + if _, err := db.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, auth, "codex", storeExecution(t, st)) + handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st)) if err != nil { t.Fatal(err) } @@ -343,13 +343,13 @@ func TestNoneSessionRetryAfterDeploymentDefaultChanges(t *testing.T) { t.Fatal("retry changed the frozen provider", err) } var revision uuid.UUID - if err := 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 := 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 { t.Fatal("API retry changed frozen revision", err) } } func TestDeploymentProviderResolutionPairsRevisionDuringReplacement(t *testing.T) { - st, pool := store.NewManagedTestStore(t) + st, db := newManagedTestStoreDB(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()}) @@ -373,7 +373,7 @@ func TestDeploymentProviderResolutionPairsRevisionDuringReplacement(t *testing.T _, err = st.SetDeploymentModelProvider(admin, harness, v1.ModelConfigurationInput{ModelProvider: replacement, Model: "fixture"}) return snapshot, err } - handler, err := publicHandler(t, st, auth, "codex", storeExecution(t, st), modelProviderDefaults(st, resolver)) + handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), modelProviderDefaults(st, resolver)) if err != nil { t.Fatal(err) } @@ -388,7 +388,7 @@ func TestDeploymentProviderResolutionPairsRevisionDuringReplacement(t *testing.T t.Fatalf("creation failed: %d %s", w.Code, w.Body) } var revision uuid.UUID - if err = 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 = 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 { t.Fatal("tuple revision changed", err) } frozen, err := st.SessionModelExecution(t.Context(), tenant, session.ID) diff --git a/services/core/internal/store/dispatch_test.go b/services/core/internal/store/dispatch_test.go index f9434316..b1b071a3 100644 --- a/services/core/internal/store/dispatch_test.go +++ b/services/core/internal/store/dispatch_test.go @@ -27,7 +27,9 @@ type dispatchHarness struct { admissions map[string]fixtureAdmission t *testing.T s *store.Store + db fixtureDB d *execution.Dispatcher + lease execution.Ownership // held by tests that run execution operations without a Worker tenant string session store.Session device store.ExecutionDevice @@ -45,8 +47,8 @@ func newDispatchHarness(t *testing.T) *dispatchHarness { func newDispatchHarnessForSession(t *testing.T, configuration []byte, local bool) *dispatchHarness { t.Helper() - s, _ := store.NewModelTestStore(t) - h := &dispatchHarness{t: t, s: s, tenant: uuid.NewString(), environments: map[string]*dispatchHarness{}} + s, db := newModelTestStoreDB(t) + h := &dispatchHarness{t: t, s: s, db: db, tenant: uuid.NewString(), environments: map[string]*dispatchHarness{}} ctx := context.Background() var err error h.session, err = s.CreateSession(ctx, h.tenant, store.WithFixtureModelProvider(store.CreateSessionInput{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: "session", Configuration: configuration})) @@ -223,9 +225,9 @@ func TestExecutionDispatchSteeringAndNativeContinuity(t *testing.T) { if outcome.AppliedThrough != second.Sequence || outcome.Done.Usage.InputTokens != 7 { t.Fatalf("missing result: %+v", outcome) } - newStore, pool := store.NewTestStore(t) - defer pool.Close() - h.s = newStore + newStore, db := newTestStoreDB(t) + defer db.pool.Close() + h.s, h.db = newStore, db h.d.Store = newStore bound, err := newStore.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || bound.NativeSessionID != "native-thread-1" { diff --git a/services/core/internal/store/environment_admission_test.go b/services/core/internal/store/environment_admission_test.go index 6e45ffe4..a8fb049d 100644 --- a/services/core/internal/store/environment_admission_test.go +++ b/services/core/internal/store/environment_admission_test.go @@ -23,22 +23,20 @@ func newEnvironmentAdmission(t *testing.T) (*dispatchHarness, *execution.Worker) t.Helper() h := newDispatchHarness(t) enableWorkerEnvironment(t, h) - worker, err := execution.StartWorker(t.Context(), h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), h.db, h.d) t.Cleanup(func() { ctx, cancel := context.WithCancel(context.Background()) cancel() _ = worker.Run(ctx) }) - h.session, err = worker.CreateSession(t.Context(), h.tenant, store.WithFixtureModelProvider(store.CreateSessionInput{ + session, err := worker.CreateSession(t.Context(), h.tenant, store.WithFixtureModelProvider(store.CreateSessionInput{ Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"model":"test-model"},"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`), })) if err != nil { t.Fatal(err) } + h.session = session h = connectFixtureRuntime(t, h, h.session) return h, worker } diff --git a/services/core/internal/store/environment_claim_worker_test.go b/services/core/internal/store/environment_claim_worker_test.go index 9fd58223..79b49742 100644 --- a/services/core/internal/store/environment_claim_worker_test.go +++ b/services/core/internal/store/environment_claim_worker_test.go @@ -14,13 +14,10 @@ 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, pool := store.NewTestStore(t) + s, db := newTestStoreDB(t) tenant, pending := newEnvironmentExpiryReservation(t, s) - writer, err := store.NewExecution(t.Context(), s) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = writer.CloseExecution(context.Background()) }) + owner := executionOwner(t, db, s) + writer := owner.Store got, err := writer.PromoteEnvironmentInput(t.Context(), tenant, pending.SessionID, pending.ID) if err != nil || len(got.Receipts) != 1 || got.Receipts[0].Replayed { t.Fatal(got, err) @@ -39,18 +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 := observeExecutionLeaseRelease(t, pool) - if err := writer.CloseExecution(t.Context()); err != nil { + awaitRelease := observeExecutionLeaseRelease(t, db.pool) + if err := owner.Lease.Close(t.Context()); err != nil { t.Fatal(err) } awaitRelease() - restarted, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()}) - if err != nil { - t.Fatal(err) - } + restarted := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()}) stopped, cancel := context.WithCancel(t.Context()) cancel() - awaitRelease = observeExecutionLeaseRelease(t, pool) + awaitRelease = observeExecutionLeaseRelease(t, db.pool) if err := restarted.Run(stopped); !errors.Is(err, context.Canceled) { t.Fatal(err) } @@ -63,18 +57,14 @@ func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { t.Fatal("restart failed to settle the original claim", turn, err) } var turns, inputs, queued int - err = pool.QueryRow(t.Context(), `SELECT + err = db.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, err := store.NewExecution(t.Context(), s) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = successor.CloseExecution(context.Background()) }) + successor := executionOwner(t, db, s).Store retry, err := successor.PromoteEnvironmentInput(t.Context(), tenant, pending.SessionID, pending.ID) if deleted { if !errors.Is(err, store.ErrNotFound) { diff --git a/services/core/internal/store/environment_connection_worker_test.go b/services/core/internal/store/environment_connection_worker_test.go index 29991aae..c7c9ed8a 100644 --- a/services/core/internal/store/environment_connection_worker_test.go +++ b/services/core/internal/store/environment_connection_worker_test.go @@ -13,7 +13,7 @@ import ( ) func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { - s, pool := store.NewTestStore(t) + s, db := newTestStoreDB(t) tenant := uuid.NewString() session, err := s.CreateSession(t.Context(), tenant, store.CreateSessionInput{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: "connection-worker", Configuration: []byte(`{"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`)}) if err != nil { @@ -23,10 +23,8 @@ func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { if err != nil { t.Fatal(err) } - writer, err := store.NewExecution(t.Context(), s) - if err != nil { - t.Fatal(err) - } + owner := executionOwner(t, db, s) + writer := owner.Store generation := uuid.NewString() if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { t.Fatal(err) @@ -34,16 +32,12 @@ func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err != nil { t.Fatal(err) } - awaitRelease := observeExecutionLeaseRelease(t, pool) - if err := writer.CloseExecution(t.Context()); err != nil { + awaitRelease := observeExecutionLeaseRelease(t, db.pool) + if err := owner.Lease.Close(t.Context()); err != nil { t.Fatal(err) } awaitRelease() - dispatcher := &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()} - worker, err := execution.StartWorker(t.Context(), dispatcher) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()}) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) exited := make(chan struct{}) diff --git a/services/core/internal/store/environment_connections_test.go b/services/core/internal/store/environment_connections_test.go index a7ba2872..f4d37558 100644 --- a/services/core/internal/store/environment_connections_test.go +++ b/services/core/internal/store/environment_connections_test.go @@ -168,7 +168,7 @@ func TestEnvironmentConnectionRequiresOwnerAndRollsBackWithEvent(t *testing.T) { if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err != nil { t.Fatal(err) } - if err := writer.CloseExecution(t.Context()); err != nil { + if err := writer.lease.Close(t.Context()); err != nil { t.Fatal(err) } before = connectionSnapshot(t, pool, environment.ID) diff --git a/services/core/internal/store/environment_device_test.go b/services/core/internal/store/environment_device_test.go index bfb67ff8..9ffa02aa 100644 --- a/services/core/internal/store/environment_device_test.go +++ b/services/core/internal/store/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.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) select { case frame := <-generalFrames: t.Fatal("general device received self-hosted work", frame.Type) diff --git a/services/core/internal/store/environment_directory_test.go b/services/core/internal/store/environment_directory_test.go index f28339ae..fa49fc0b 100644 --- a/services/core/internal/store/environment_directory_test.go +++ b/services/core/internal/store/environment_directory_test.go @@ -34,10 +34,7 @@ func directoryWorker(t *testing.T) (*dispatchHarness, *execution.Worker, store.E info, _, _ := peer.AgentKindStatus("codex") return info.Capabilities.WorkspaceReadPreparation }) - w, err := execution.StartWorker(t.Context(), h.d) - if err != nil { - t.Fatal(err) - } + w := startWorker(t, t.Context(), h.db, h.d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- w.Run(ctx) }() diff --git a/services/core/internal/store/environment_executor_credentials_test.go b/services/core/internal/store/environment_executor_credentials_test.go index ab67e372..19f61aa1 100644 --- a/services/core/internal/store/environment_executor_credentials_test.go +++ b/services/core/internal/store/environment_executor_credentials_test.go @@ -122,11 +122,7 @@ func TestEnvironmentExecutorConcurrentIssueAndDeletion(t *testing.T) { t.Fatal(err) } // Provisioning remains control-plane work while the execution owner is active. - writer, err := NewExecution(ctx, s) - if err != nil { - t.Fatal(err) - } - defer writer.CloseExecution(ctx) + executionWriter(t, s) const attempts = 8 var wg sync.WaitGroup tokens := make(chan string, attempts) diff --git a/services/core/internal/store/environment_expiry_dispatch_test.go b/services/core/internal/store/environment_expiry_dispatch_test.go index efb64445..7f36117d 100644 --- a/services/core/internal/store/environment_expiry_dispatch_test.go +++ b/services/core/internal/store/environment_expiry_dispatch_test.go @@ -19,7 +19,7 @@ func TestWorkerEnvironmentExpiryAtFullExecutionCapacity(t *testing.T) { h := newDispatchHarness(t) _, pool := store.NewTestStore(t) enableEnvironmentExpiryDispatch(h) - worker, stop := startEnvironmentExpiryWorker(t, h.d) + worker, stop := startEnvironmentExpiryWorker(t, h.db, h.d) var requests []proto.Envelope var sessions []store.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.d) + worker, stop := startEnvironmentExpiryWorker(t, h.db, h.d) h.session = publicSession(t, h, "unrelated") receipt, err := worker.SubmitInputs(t.Context(), h.tenant, h.session.ID, "work", []store.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"make normal progress"}`)}}) if err != nil { diff --git a/services/core/internal/store/environment_expiry_worker_test.go b/services/core/internal/store/environment_expiry_worker_test.go index 51b4cce5..48504057 100644 --- a/services/core/internal/store/environment_expiry_worker_test.go +++ b/services/core/internal/store/environment_expiry_worker_test.go @@ -39,12 +39,9 @@ func makeEnvironmentExpiryDue(t *testing.T, pool *pgxpool.Pool, pending *store.E } } -func startEnvironmentExpiryWorker(t *testing.T, d *execution.Dispatcher) (*execution.Worker, func()) { +func startEnvironmentExpiryWorker(t *testing.T, db fixtureDB, d *execution.Dispatcher) (*execution.Worker, func()) { t.Helper() - worker, err := execution.StartWorker(t.Context(), d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -102,22 +99,22 @@ func assertEnvironmentExpiryHasNoHistory(t *testing.T, pool *pgxpool.Pool, sessi } func TestWorkerEnvironmentExpiryWithoutDevicesAndAfterRestart(t *testing.T) { - s, pool := store.NewTestStore(t) + s, db := newTestStoreDB(t) dueTenant, due := newEnvironmentExpiryReservation(t, s) futureTenant, future := newEnvironmentExpiryReservation(t, s) - makeEnvironmentExpiryDue(t, pool, &due) + makeEnvironmentExpiryDue(t, db.pool, &due) d := &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()} - _, stop := startEnvironmentExpiryWorker(t, d) + _, stop := startEnvironmentExpiryWorker(t, db, d) waitEnvironmentExpiry(t, s, dueTenant, due) got, err := s.GetEnvironmentInputReservation(t.Context(), futureTenant, future.SessionID, future.ID) if err != nil || got.State != store.EnvironmentInputPending || !got.Deadline.Equal(future.Deadline) { t.Fatal("future input changed", got, err) } stop() - makeEnvironmentExpiryDue(t, pool, &future) - _, stop = startEnvironmentExpiryWorker(t, d) + makeEnvironmentExpiryDue(t, db.pool, &future) + _, stop = startEnvironmentExpiryWorker(t, db, d) waitEnvironmentExpiry(t, s, futureTenant, future) stop() - assertEnvironmentExpiryHasNoHistory(t, pool, due.SessionID) - assertEnvironmentExpiryHasNoHistory(t, pool, future.SessionID) + assertEnvironmentExpiryHasNoHistory(t, db.pool, due.SessionID) + assertEnvironmentExpiryHasNoHistory(t, db.pool, future.SessionID) } diff --git a/services/core/internal/store/environment_file_write_semantics_public_test.go b/services/core/internal/store/environment_file_write_semantics_public_test.go index aa111057..9b74a753 100644 --- a/services/core/internal/store/environment_file_write_semantics_public_test.go +++ b/services/core/internal/store/environment_file_write_semantics_public_test.go @@ -59,7 +59,7 @@ 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, auth, "codex", workerExecution(w)) + handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(w)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_initial_input_test.go b/services/core/internal/store/environment_initial_input_test.go index d2d3d90b..55c1838b 100644 --- a/services/core/internal/store/environment_initial_input_test.go +++ b/services/core/internal/store/environment_initial_input_test.go @@ -227,7 +227,7 @@ func TestEnvironmentInitialInputExpiryHasNoTurnAndCannotReplay(t *testing.T) { if err != nil || len(events) != expectedEvents || events[len(events)-1].Event.Type != "agent.session.failed" || events[len(events)-1].Turn != nil || events[len(events)-1].EnvironmentInputActivity.Status != "failed" { t.Fatal("missing pre-Turn failure snapshot", events, err) } - if err := writer.CloseExecution(t.Context()); err != nil { + if err := writer.lease.Close(t.Context()); err != nil { t.Fatal(err) } pool.Close() diff --git a/services/core/internal/store/environment_initial_public_test.go b/services/core/internal/store/environment_initial_public_test.go index 249af4af..6147b8db 100644 --- a/services/core/internal/store/environment_initial_public_test.go +++ b/services/core/internal/store/environment_initial_public_test.go @@ -21,7 +21,7 @@ func TestEnvironmentInitialFailureOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, pool := store.NewTestStore(t) + s, db := newTestStoreDB(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,18 +36,16 @@ func TestEnvironmentInitialFailureOfficialClient(t *testing.T) { if err != nil { t.Fatal(err) } - writer, err := store.NewExecution(t.Context(), s) - if err != nil { - t.Fatal(err) - } + owner := executionOwner(t, db, s) + writer := owner.Store t.Cleanup(func() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if err := writer.CloseExecution(ctx); err != nil { + if err := owner.Lease.Close(ctx); err != nil { t.Error(err) } }) - handler, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) + handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } @@ -92,7 +90,7 @@ func TestEnvironmentInitialFailureOfficialClient(t *testing.T) { } } var reservation string - if err := 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 := 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 { t.Fatal(err) } if result, err := writer.ExpireEnvironmentInput(t.Context(), tenant, session.ID, reservation); err != nil || result.State != store.EnvironmentInputExpired { diff --git a/services/core/internal/store/environment_initialization.go b/services/core/internal/store/environment_initialization.go index 6964e321..5aeed5e6 100644 --- a/services/core/internal/store/environment_initialization.go +++ b/services/core/internal/store/environment_initialization.go @@ -16,7 +16,7 @@ type EnvironmentInitialization struct { } func (s *Store) ListEnvironmentInitializations(ctx context.Context, after string) ([]EnvironmentInitialization, error) { - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return nil, err } id := pgtype.UUID{Valid: true} diff --git a/services/core/internal/store/environment_initialization_test.go b/services/core/internal/store/environment_initialization_test.go index 5f8a3c36..9f5aba92 100644 --- a/services/core/internal/store/environment_initialization_test.go +++ b/services/core/internal/store/environment_initialization_test.go @@ -52,7 +52,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} principal := store.FixtureExecutorPrincipal(t, s, uuid.NewString()) session, err := s.CreateSession(t.Context(), principal.TenantID, store.CreateSessionInput{ Creator: principal.Subject(), Engine: "codex", IdempotencyKey: uuid.NewString(), @@ -79,10 +79,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(s), Registry: registry}) server := httptest.NewServer(http.HandlerFunc(handler.WS)) defer server.Close() - worker, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: registry}) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: registry}) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -171,7 +168,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, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) principal := store.FixtureExecutorPrincipal(t, s, uuid.NewString()) create := func() store.EnvironmentInitialization { t.Helper() @@ -198,11 +195,7 @@ func TestEnvironmentInitializationRevocationBeforeClaim(t *testing.T) { return store.EnvironmentInitialization{EnvironmentID: environment.ID, SessionID: session.ID, TenantID: principal.TenantID, DeviceID: enrolled.DeviceID, State: "pending", Engine: "codex"} } revoked, other := create(), create() - owned, err := store.NewExecution(t.Context(), s) - if err != nil { - t.Fatal(err) - } - defer owned.CloseExecution(context.Background()) + owned := executionOwner(t, db, s).Store if err := s.RevokeDevice(t.Context(), principal.TenantID, revoked.DeviceID); err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_input_activity_test.go b/services/core/internal/store/environment_input_activity_test.go index b1509e9b..62490b22 100644 --- a/services/core/internal/store/environment_input_activity_test.go +++ b/services/core/internal/store/environment_input_activity_test.go @@ -215,7 +215,7 @@ func TestEnvironmentInputActivityRecoversWaitingActionAndHidesDeletion(t *testin t.Fatal(err) } requireEnvironmentInputActivity(t, s, tenant, session.ID, "idle", "") - if err := old.CloseExecution(t.Context()); err != nil { + if err := old.lease.Close(t.Context()); err != nil { t.Fatal(err) } next := executionWriter(t, s) diff --git a/services/core/internal/store/environment_input_migration_test.go b/services/core/internal/store/environment_input_migration_test.go index e670d627..68c58743 100644 --- a/services/core/internal/store/environment_input_migration_test.go +++ b/services/core/internal/store/environment_input_migration_test.go @@ -99,7 +99,7 @@ func TestEnvironmentInputPromotionUsesCurrentExecutionWriter(t *testing.T) { t.Fatal("pooled Store promoted input without execution ownership") } closed := executionWriter(t, s) - if err := closed.CloseExecution(t.Context()); err != nil { + if err := closed.lease.Close(t.Context()); err != nil { t.Fatal(err) } if _, err := closed.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID); err == nil { diff --git a/services/core/internal/store/environment_inputs_test.go b/services/core/internal/store/environment_inputs_test.go index 55cd01b1..30d82df8 100644 --- a/services/core/internal/store/environment_inputs_test.go +++ b/services/core/internal/store/environment_inputs_test.go @@ -165,7 +165,7 @@ func TestEnvironmentInputReservationPromotionAndDirectRetries(t *testing.T) { if err != nil || len(retry) != 2 || !retry[0].Replayed || retry[0].Sequence != promoted.Receipts[0].Sequence { t.Fatal("direct retry after promotion", retry, err) } - if err := writer.CloseExecution(ctx); err != nil { + if err := writer.lease.Close(ctx); err != nil { t.Fatal(err) } pool.Close() diff --git a/services/core/internal/store/environment_mcp_public_test.go b/services/core/internal/store/environment_mcp_public_test.go index ba21a2aa..35a0c47d 100644 --- a/services/core/internal/store/environment_mcp_public_test.go +++ b/services/core/internal/store/environment_mcp_public_test.go @@ -15,9 +15,9 @@ import ( func TestPublicEnvironmentMCPUsesAttachedVaultSelection(t *testing.T) { for _, kind := range []string{"codex", "claude_sdk", "mcode"} { t.Run(kind, func(t *testing.T) { - s, pool, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) + s, db, 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, auth, kind, workerExecution(&execution.Worker{}), executorURL("https://executor.example")) + handler, err := publicHandler(t, s, db, auth, kind, workerExecution(&execution.Worker{}), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } @@ -48,7 +48,7 @@ func TestPublicEnvironmentMCPUsesAttachedVaultSelection(t *testing.T) { } } } - assertSelfHostedMCPRejectionHasNoWrites(t, pool, tenant) + assertSelfHostedMCPRejectionHasNoWrites(t, db.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/internal/store/environment_retrieve_public_test.go b/services/core/internal/store/environment_retrieve_public_test.go index 892a2038..cc9036be 100644 --- a/services/core/internal/store/environment_retrieve_public_test.go +++ b/services/core/internal/store/environment_retrieve_public_test.go @@ -22,7 +22,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, pool := store.NewModelTestStore(t) + s, db := newModelTestStoreDB(t) tenant, foreignTenant := uuid.NewString(), uuid.NewString() principal := store.FixtureExecutorPrincipal(t, s, tenant) token, peer, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() @@ -48,7 +48,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { } } }() - handler, err := publicHandler(t, s, auth, "codex", storeExecution(t, s), executorURL("https://private-registry.example")) + handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://private-registry.example")) if err != nil { t.Fatal(err) } @@ -85,9 +85,9 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { } revoked = true server.Close() - pool.Close() - reopened, reopenedPool := store.NewModelTestStore(t) - handler, err = publicHandler(t, reopened, auth, "codex") + db.pool.Close() + reopened, reopenedDB := newModelTestStoreDB(t) + handler, err = publicHandler(t, reopened, reopenedDB, auth, "codex") if err != nil { t.Fatal(err) } @@ -103,7 +103,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { t.Fatal("public retrieval changed durable Environment state", err) } var history int - if err := reopenedPool.QueryRow(t.Context(), `SELECT + if err := reopenedDB.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/internal/store/environment_runtime_fixture_test.go b/services/core/internal/store/environment_runtime_fixture_test.go index 5d2ac39c..ba78d9e3 100644 --- a/services/core/internal/store/environment_runtime_fixture_test.go +++ b/services/core/internal/store/environment_runtime_fixture_test.go @@ -38,7 +38,7 @@ func enrollFixtureSession(t *testing.T, s *store.Store, tenant string, session s func connectFixtureRuntime(t *testing.T, h *dispatchHarness, session store.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, d: h.d, tenant: h.tenant, session: session, registry: h.registry, url: h.url, + other := &dispatchHarness{t: h.t, s: h.s, db: h.db, lease: h.lease, 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.tenant, session) u, err := url.Parse(h.url) diff --git a/services/core/internal/store/environment_worker_scan_test.go b/services/core/internal/store/environment_worker_scan_test.go index ee7b76cc..f07eddbe 100644 --- a/services/core/internal/store/environment_worker_scan_test.go +++ b/services/core/internal/store/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.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, 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.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) barrier := nextWorkerFrame(t, frames, testExecutionRequest) if barrier.ID != receipt.TurnID { t.Fatal("unexpected page barrier") diff --git a/services/core/internal/store/environment_worker_test.go b/services/core/internal/store/environment_worker_test.go index 37133ef7..86742602 100644 --- a/services/core/internal/store/environment_worker_test.go +++ b/services/core/internal/store/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.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, 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.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, 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 != store.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.d) + _, stop = startEnvironmentExpiryWorker(t, h.db, 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/internal/store/environment_write_audit_test.go b/services/core/internal/store/environment_write_audit_test.go index 7de751b8..5fc200b5 100644 --- a/services/core/internal/store/environment_write_audit_test.go +++ b/services/core/internal/store/environment_write_audit_test.go @@ -21,7 +21,7 @@ func TestEnvironmentUploadWriteAuditSurvivesRequestAndLeaseContext(t *testing.T) cancel() sessionAuditCount(t, f.s, f.tenant, 0) // Neither the settlement caller nor the newly acquired execution lease owns the request context. - if err := f.writer.CloseExecution(context.Background()); err != nil { + if err := f.writer.lease.Close(context.Background()); err != nil { t.Fatal(err) } next := executionWriter(t, f.s) diff --git a/services/core/internal/store/execution.go b/services/core/internal/store/execution.go index bf56baf9..c820c55b 100644 --- a/services/core/internal/store/execution.go +++ b/services/core/internal/store/execution.go @@ -17,19 +17,14 @@ type transactor interface { Transaction(context.Context, func(context.Context, pgx.Tx) error) error } -// NewExecution takes the database's execution lease on a dedicated connection -// and returns the execution writer built on it. The writer's Session and -// execution-only transactions run on the leased connection; reads keep the pool. -// Keep public admission on the pooled Store. Losing or closing the lease never -// falls back to a pooled writer. It fails when another service owns the database. -func NewExecution(ctx context.Context, s *Store) (*Store, error) { - lease, err := pgunit.AcquireLease(ctx, s.pool) - if err != nil { - return nil, err - } +// NewExecution returns the execution writer built on lease, which the caller +// acquired and closes. The writer's Session and execution-only transactions run +// on the leased connection; reads keep the pool. Keep public admission on the +// pooled Store. Losing or closing the lease never falls back to a pooled writer. +func NewExecution(s *Store, lease *pgunit.Lease) *Store { writer := *s writer.writer, writer.lease = lease, lease - return &writer, nil + return &writer } // checkExecutionAuthority only validates. The connection was fixed when the @@ -41,28 +36,11 @@ func (s *Store) checkExecutionAuthority() error { return nil } -// CheckExecutionOwnership validates the current writer before external preparation. -func (s *Store) CheckExecutionOwnership(ctx context.Context) error { +// checkExecutionOwnership confirms, before an execution-only pooled read, that +// the borrowed lease still owns the database. +func (s *Store) checkExecutionOwnership(ctx context.Context) error { if err := s.checkExecutionAuthority(); err != nil { return err } return s.lease.CheckOwnership(ctx) } - -// CancelExecutionOperations cancels coordinator-owned contexts between leased -// operations; see pgunit.Lease.CancelOperations for the constraints on cancel. -func (s *Store) CancelExecutionOperations(ctx context.Context, cancel context.CancelFunc) error { - if err := s.checkExecutionAuthority(); err != nil { - return err - } - return s.lease.CancelOperations(ctx, cancel) -} - -// CloseExecution releases the execution lease and waits for connection cleanup -// within ctx. A later call resumes that wait. The writer cannot write afterwards. -func (s *Store) CloseExecution(ctx context.Context) error { - if err := s.checkExecutionAuthority(); err != nil { - return err - } - return s.lease.Close(ctx) -} diff --git a/services/core/internal/store/execution_test.go b/services/core/internal/store/execution_test.go index 32ff3260..1382472e 100644 --- a/services/core/internal/store/execution_test.go +++ b/services/core/internal/store/execution_test.go @@ -13,19 +13,21 @@ import ( "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" ) -// executionWriter builds the execution writer on the shared test database. -// Tests in this package run sequentially, so one writer at a time owns it. +// executionWriter acquires the execution lease on the shared test database and +// builds s's execution writer on it, as cmd/server does. Tests in this package +// run sequentially, so one writer at a time owns it; writer.lease releases it. func executionWriter(t *testing.T, s *Store) *Store { t.Helper() - writer, err := NewExecution(t.Context(), s) + lease, err := pgunit.AcquireLease(t.Context(), s.pool) if err != nil { t.Fatal(err) } - t.Cleanup(func() { _ = writer.CloseExecution(context.Background()) }) - return writer + t.Cleanup(func() { _ = lease.Close(context.Background()) }) + return NewExecution(s, lease) } // executionOwnerPID finds the backend holding this database's execution lease, @@ -105,7 +107,7 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { mustReject("reconciliation", err) _, err = writer.ExpireEnvironmentInputs(t.Context()) mustReject("input expiry", err) - mustReject("ownership check", writer.CheckExecutionOwnership(t.Context())) + mustReject("ownership check", writer.lease.CheckOwnership(t.Context())) after, err := s.GetSession(t.Context(), tenant, active.ID) if err != nil || !reflect.DeepEqual(before, after) { t.Fatal("stale state persisted", after, err) @@ -131,7 +133,7 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { } // Public admission remains usable with a dead owner connection. submitMessage(t, s, tenant, queued.ID, "additional") - if err = successor.CheckExecutionOwnership(t.Context()); err != nil { + if err = successor.lease.CheckOwnership(t.Context()); err != nil { t.Fatal(err) } if _, err = successor.CompleteExecution(t.Context(), tenant, active.ID, input.TurnID, TurnCompleted, json.RawMessage(`{"done":{"content":"accepted"}}`), "successor-native", input.Sequence); err != nil { @@ -145,7 +147,7 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { if !errors.Is(err, ErrTurnConflict) { t.Fatal("terminal CAS changed", err) } - if err = successor.CloseExecution(t.Context()); err != nil { + if err = successor.lease.Close(t.Context()); err != nil { t.Fatal(err) } mustReject("closed writer", successor.BindSessionDevice(t.Context(), tenant, queued.ID, host.ID)) @@ -184,7 +186,7 @@ func TestExecutionWriterSerializesWritesOnItsLease(t *testing.T) { } results <- err }) - group.Go(func() { results <- writer.CheckExecutionOwnership(t.Context()) }) + group.Go(func() { results <- writer.lease.CheckOwnership(t.Context()) }) } group.Wait() close(results) @@ -230,15 +232,11 @@ func TestPooledStoreHasNoExecutionAuthority(t *testing.T) { if err != nil { t.Fatal(err) } - operation, cancel := context.WithCancel(t.Context()) - defer cancel() _, archiveErr := s.ArchiveManagedSession(t.Context(), tenant, session.ID, 0) _, expiryErr := s.ExpireEnvironmentInputs(t.Context()) subagent := []ExecutionEvent{{Kind: "delta", Payload: json.RawMessage(`{"delta":"visible"}`)}, subagentIdentityEvent("child", "root", 100)} for name, err := range map[string]error{ - "ownership check": s.CheckExecutionOwnership(t.Context()), - "cancellation": s.CancelExecutionOperations(t.Context(), cancel), - "close": s.CloseExecution(t.Context()), + "ownership check": s.checkExecutionOwnership(t.Context()), "archive": archiveErr, "input expiry": expiryErr, "deployment": s.ConfigureRuntimeDeployment(t.Context(), nil), @@ -251,9 +249,6 @@ func TestPooledStoreHasNoExecutionAuthority(t *testing.T) { t.Fatalf("pooled Store ran %s: %v", name, err) } } - if operation.Err() != nil { - t.Fatal("pooled Store ran the cancellation callback") - } after, err := s.GetSession(t.Context(), tenant, session.ID) if err != nil || !reflect.DeepEqual(before, after) { t.Fatal("rejected execution operation changed the Session", after, err) diff --git a/services/core/internal/store/file_resource_semantics_public_test.go b/services/core/internal/store/file_resource_semantics_public_test.go index 9a7c4bf7..2007062e 100644 --- a/services/core/internal/store/file_resource_semantics_public_test.go +++ b/services/core/internal/store/file_resource_semantics_public_test.go @@ -33,8 +33,8 @@ func TestFileResourceSemanticsOfficialClientPostgres(t *testing.T) { }) newServer := func() *httptest.Server { t.Helper() - s := store.NewWithCredentialCipher(pool, cipher) - h, err := publicHandler(t, s, auth, "codex") + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + h, err := publicHandler(t, s, db, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/fixture_db_test.go b/services/core/internal/store/fixture_db_test.go new file mode 100644 index 00000000..ac144563 --- /dev/null +++ b/services/core/internal/store/fixture_db_test.go @@ -0,0 +1,72 @@ +package store_test + +import ( + "context" + "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/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" +) + +// fixtureDB is the database and credential key that built the test's Store. +// Cutovers build their adapters from it; add fields here, never parameters. +type fixtureDB struct { + pool *pgxpool.Pool + cipher *credentialcrypto.Cipher // nil for a keyless Store +} + +// newTestStoreDB is store.NewTestStore with the fixtureDB that built it. +func newTestStoreDB(t *testing.T) (*store.Store, fixtureDB) { + s, pool := store.NewTestStore(t) + return s, fixtureDB{pool: pool} +} + +// newModelTestStoreDB is store.NewModelTestStore with the fixtureDB that built it. +func newModelTestStoreDB(t *testing.T) (*store.Store, fixtureDB) { + s, pool := store.NewModelTestStore(t) + return s, fixtureDB{pool: pool, cipher: store.FixtureCipher()} +} + +// newManagedTestStoreDB is store.NewManagedTestStore with the fixtureDB that built it. +func newManagedTestStoreDB(t *testing.T) (*store.Store, fixtureDB) { + s, pool := store.NewManagedTestStore(t) + return s, fixtureDB{pool: pool, cipher: store.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. +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 + } + return execution.StartWorker(ctx, dispatcher, execution.Owner{Lease: lease, Store: store.NewExecution(dispatcher.Store, lease)}) +} + +// executionOwner acquires the execution lease on db and builds s's execution +// writer on it, for tests that run execution operations without a Worker. The +// lease closes when the test ends. +func executionOwner(t testing.TB, db fixtureDB, s *store.Store) 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()) }) + return execution.Owner{Lease: lease, Store: store.NewExecution(s, lease)} +} diff --git a/services/core/internal/store/function_images_native_test.go b/services/core/internal/store/function_images_native_test.go index a9f9b32d..3fb73082 100644 --- a/services/core/internal/store/function_images_native_test.go +++ b/services/core/internal/store/function_images_native_test.go @@ -10,7 +10,6 @@ import ( "testing" "time" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/google/uuid" ) @@ -32,10 +31,7 @@ func TestNativeFunctionImagePublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -51,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, auth, kind, workerExecution(worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(h.s, model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(h.s, model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/function_inputs_public_test.go b/services/core/internal/store/function_inputs_public_test.go index df973ba0..9613a3b3 100644 --- a/services/core/internal/store/function_inputs_public_test.go +++ b/services/core/internal/store/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, _ := store.NewTestStore(t) + s, db := newTestStoreDB(t) ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second) defer cancel() tenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() @@ -46,7 +46,7 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { } } 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, auth, "codex", storeExecution(t, s)) + handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/function_public_native_test.go b/services/core/internal/store/function_public_native_test.go index 41b7523f..238cfaa9 100644 --- a/services/core/internal/store/function_public_native_test.go +++ b/services/core/internal/store/function_public_native_test.go @@ -11,7 +11,6 @@ import ( "time" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/google/uuid" ) @@ -61,10 +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, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) ctx, cancel := context.WithCancel(ctx) t.Cleanup(cancel) done := make(chan error, 1) @@ -79,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, auth, "codex", workerExecution(worker), nativeDeploymentDefaults(h.s, "gpt-5.5", provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(worker), nativeDeploymentDefaults(h.s, "gpt-5.5", provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/function_state_public_test.go b/services/core/internal/store/function_state_public_test.go index 5e63aa8b..561c5f76 100644 --- a/services/core/internal/store/function_state_public_test.go +++ b/services/core/internal/store/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, _ := store.NewTestStore(t) + s, db := newTestStoreDB(t) ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second) defer cancel() tenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() @@ -45,7 +45,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, auth, "codex") + handler, err := publicHandler(t, s, db, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/function_worker_test.go b/services/core/internal/store/function_worker_test.go index f964c9c4..ad31a235 100644 --- a/services/core/internal/store/function_worker_test.go +++ b/services/core/internal/store/function_worker_test.go @@ -65,10 +65,7 @@ func TestWorkerWaitsForToolCapabilities(t *testing.T) { input := h.message("queued", "Look up ticket") ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -127,7 +124,7 @@ func mcpBearerWorkerConfiguration(t *testing.T, h *dispatchHarness) (string, str if err != nil { t.Fatal(err) } - h.s = store.NewWithCredentialCipher(pool, cipher) + h.s, h.db = store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} h.d.Store = h.s vault, err := h.s.CreateVault(t.Context(), h.tenant, store.CreateVaultInput{}) if err != nil { diff --git a/services/core/internal/store/harness_onboarding_test.go b/services/core/internal/store/harness_onboarding_test.go index 209944f6..5ec4ee61 100644 --- a/services/core/internal/store/harness_onboarding_test.go +++ b/services/core/internal/store/harness_onboarding_test.go @@ -49,10 +49,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { started, write, declaration := startOnboardingPeer(t, h) ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) stopped := make(chan error, 1) go func() { stopped <- worker.Run(ctx) }() defer func() { @@ -65,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, auth, "fixture_harness", workerExecution(worker), withPolicy(policy)) + handler, err := publicHandler(t, h.s, h.db, auth, "fixture_harness", workerExecution(worker), withPolicy(policy)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/hosted_initialization_failure_public_test.go b/services/core/internal/store/hosted_initialization_failure_public_test.go index 2a334967..1c2d6117 100644 --- a/services/core/internal/store/hosted_initialization_failure_public_test.go +++ b/services/core/internal/store/hosted_initialization_failure_public_test.go @@ -98,14 +98,14 @@ func (p *hostedFailureProvider) prepare(request proto.RuntimePreparePayload, _ [ return completedInitialization(request, nil) } -func hostedFailureStore(t *testing.T) *store.Store { +func hostedFailureStore(t *testing.T) (*store.Store, fixtureDB) { t.Helper() _, pool := store.NewManagedTestStore(t) cipher, err := credentialcrypto.New(bytes.Repeat([]byte{7}, 32)) if err != nil { t.Fatal(err) } - return store.NewWithCredentialCipher(pool, cipher) + return store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} } func hostedFailureSession(t *testing.T, s *store.Store, tenant string, input store.CreateSessionInput) (store.Session, store.Environment) { @@ -126,10 +126,10 @@ func hostedFailureSession(t *testing.T, s *store.Store, tenant string, input sto return session, environment } -func failHostedInitialization(t *testing.T, s *store.Store, tenant string, environment store.Environment, p *hostedFailureProvider) { +func failHostedInitialization(t *testing.T, s *store.Store, db fixtureDB, tenant string, environment store.Environment, p *hostedFailureProvider) { t.Helper() key := uuid.NewString() - w, _ := managedWorkerMode(t, s, key, p, false, true) + w, _ := managedWorkerMode(t, s, db, key, p, false, true) if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err != nil { t.Fatal(err) } @@ -180,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 := hostedFailureStore(t) + s, db := hostedFailureStore(t) tenant := uuid.NewString() 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, tenant, environment, p) + failHostedInitialization(t, s, db, 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) } @@ -245,7 +245,7 @@ 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 := hostedFailureStore(t) + s, db := hostedFailureStore(t) tenant := uuid.NewString() session, environment := hostedFailureSession(t, s, tenant, store.CreateSessionInput{ Initialization: environmentconfig.Setup{Commands: []environmentconfig.SetupCommand{{Command: "exit 3"}}}, @@ -253,7 +253,7 @@ func TestHostedInitializationFailureSettlesPendingInitialInput(t *testing.T) { }) p := &hostedFailureProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, fail: "setup", result: failedInitialization(3)} - failHostedInitialization(t, s, tenant, environment, p) + failHostedInitialization(t, s, db, tenant, environment, p) read, err := 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 { @@ -279,7 +279,7 @@ 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 := hostedFailureStore(t) + s, db := hostedFailureStore(t) tenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() session, environment := hostedFailureSession(t, s, tenant, store.CreateSessionInput{ Initialization: environmentconfig.Setup{Commands: []environmentconfig.SetupCommand{{Command: "echo " + hostedFailureCanary + "; exit 3"}}}, @@ -292,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, key, p, false, true) + w, _ := managedWorkerMode(t, s, db, 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, auth, "codex", workerExecution(w)) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(w)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/initial_files_http_test.go b/services/core/internal/store/initial_files_http_test.go index d70d54ac..88c98569 100644 --- a/services/core/internal/store/initial_files_http_test.go +++ b/services/core/internal/store/initial_files_http_test.go @@ -20,11 +20,11 @@ func TestInitialFilesHTTPInlineLimitsAndRetry(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, s), managedSandboxes(t, s), fixtureDeploymentProvider(s)) + handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), managedSandboxes(t, s), fixtureDeploymentProvider(s)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/input_conflicts_public_test.go b/services/core/internal/store/input_conflicts_public_test.go index de54583e..dfb46f23 100644 --- a/services/core/internal/store/input_conflicts_public_test.go +++ b/services/core/internal/store/input_conflicts_public_test.go @@ -37,14 +37,14 @@ 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, pool := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(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, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } @@ -132,7 +132,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, pool) + digest := databaseDigest(t, db.pool) before := make([]string, len(watched)) for i, session := range watched { before[i] = read(session) @@ -146,7 +146,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { t.Errorf("%s: %d %s", tc.name, status, body) } } - if after := databaseDigest(t, pool); !reflect.DeepEqual(after, digest) { + if after := databaseDigest(t, db.pool); !reflect.DeepEqual(after, digest) { t.Error("rejected input changed the database") } for i, session := range watched { @@ -271,7 +271,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { if status, body := createSession("first"); status != http.StatusCreated { t.Fatal(status, body) } - digest := databaseDigest(t, pool) + digest := databaseDigest(t, db.pool) status, body := createSession("changed") var creation struct { Error struct { @@ -284,7 +284,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, pool); !reflect.DeepEqual(after, digest) { + if after := databaseDigest(t, db.pool); !reflect.DeepEqual(after, digest) { t.Error("creation key reuse changed the database") } } diff --git a/services/core/internal/store/list_cursor_public_test.go b/services/core/internal/store/list_cursor_public_test.go index 6a005701..7b949357 100644 --- a/services/core/internal/store/list_cursor_public_test.go +++ b/services/core/internal/store/list_cursor_public_test.go @@ -221,25 +221,21 @@ func TestListCursorErrorsPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() client := pathIDClient{t: t, server: server} - writer, err := store.NewExecution(t.Context(), s) - if err != nil { - t.Fatal(err) - } - defer func() { _ = writer.CloseExecution(t.Context()) }() + writer := executionOwner(t, db, s).Store a := seedCursorFixture(t, s, writer, client, owner, ownerTenant, "a") b := seedCursorFixture(t, s, writer, client, foreign, foreignTenant, "b") diff --git a/services/core/internal/store/list_query_public_test.go b/services/core/internal/store/list_query_public_test.go index e8faa2df..72bcf988 100644 --- a/services/core/internal/store/list_query_public_test.go +++ b/services/core/internal/store/list_query_public_test.go @@ -26,7 +26,7 @@ func TestListQueryOfficialClientPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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()}, @@ -34,10 +34,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, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s}) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s}) t.Cleanup(func() { stopped, cancel := context.WithCancel(context.Background()) cancel() @@ -45,7 +42,7 @@ func TestListQueryOfficialClientPostgres(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, auth, "codex", workerExecution(worker)) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(worker)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/local_environment_worker_test.go b/services/core/internal/store/local_environment_worker_test.go index 971d405a..53e73c0c 100644 --- a/services/core/internal/store/local_environment_worker_test.go +++ b/services/core/internal/store/local_environment_worker_test.go @@ -41,10 +41,7 @@ func localWorker(t *testing.T, scoped, execute bool) (*dispatchHarness, *executi info, _, _ := peer.AgentKindStatus("codex") return info.Capabilities.LocalEnvironment }) - w, err := execution.StartWorker(t.Context(), h.d) - if err != nil { - t.Fatal(err) - } + w := startWorker(t, t.Context(), h.db, h.d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- w.Run(ctx) }() diff --git a/services/core/internal/store/mcode_public_native_test.go b/services/core/internal/store/mcode_public_native_test.go index 51e1c9c7..f71ac1de 100644 --- a/services/core/internal/store/mcode_public_native_test.go +++ b/services/core/internal/store/mcode_public_native_test.go @@ -31,10 +31,7 @@ func TestNativeMCodePublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 12*time.Minute) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) stopped := make(chan error, 1) go func() { stopped <- worker.Run(ctx) }() defer func() { @@ -50,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, auth, "mcode", workerExecution(worker), acceptUnavailable(t), nativeDeploymentDefaults(h.s, model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "mcode", workerExecution(worker), acceptUnavailable(t), nativeDeploymentDefaults(h.s, model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/mcp_credential_selection_public_test.go b/services/core/internal/store/mcp_credential_selection_public_test.go index 62629755..39352c2c 100644 --- a/services/core/internal/store/mcp_credential_selection_public_test.go +++ b/services/core/internal/store/mcp_credential_selection_public_test.go @@ -30,13 +30,13 @@ func TestMCPCredentialSelectionPublicPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/message_image_admission_test.go b/services/core/internal/store/message_image_admission_test.go index f1202c08..9a356ac7 100644 --- a/services/core/internal/store/message_image_admission_test.go +++ b/services/core/internal/store/message_image_admission_test.go @@ -28,10 +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, err := execution.StartWorker(t.Context(), h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), h.db, 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 := store.CreateSessionInput{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: "image-create", Configuration: configuration, InitialInputs: imageAdmissionBatch()} diff --git a/services/core/internal/store/message_images_native_test.go b/services/core/internal/store/message_images_native_test.go index 31a0862d..29002b5e 100644 --- a/services/core/internal/store/message_images_native_test.go +++ b/services/core/internal/store/message_images_native_test.go @@ -32,10 +32,7 @@ func TestNativeMessageImagePublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -51,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, auth, kind, workerExecution(worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(h.s, model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(h.s, model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/model_protocol_native_test.go b/services/core/internal/store/model_protocol_native_test.go index 4172e077..1fec8ba4 100644 --- a/services/core/internal/store/model_protocol_native_test.go +++ b/services/core/internal/store/model_protocol_native_test.go @@ -12,7 +12,6 @@ import ( "time" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" @@ -62,7 +61,7 @@ func TestNativeModelProtocolPublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 15*time.Minute) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) + worker, err := startWorkerErr(ctx, h.db, h.d) if err != nil { t.Fatal("cannot start native execution worker") } @@ -79,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, auth, options.Engine, workerExecution(worker), withPolicy(h.d.Policy), modelProviderDefaults(h.s, func(context.Context, string) (*store.DeploymentModelProviderSnapshot, error) { + handler, err := publicHandler(t, h.s, h.db, auth, options.Engine, workerExecution(worker), withPolicy(h.d.Policy), modelProviderDefaults(h.s, func(context.Context, string) (*store.DeploymentModelProviderSnapshot, error) { return &store.DeploymentModelProviderSnapshot{Model: options.Model, HarnessConfig: options.HarnessConfig, Provider: &options.Provider, Revision: providerRevision}, nil })) if err != nil { diff --git a/services/core/internal/store/native_public_execution_test.go b/services/core/internal/store/native_public_execution_test.go index de2aaf95..2eed0cad 100644 --- a/services/core/internal/store/native_public_execution_test.go +++ b/services/core/internal/store/native_public_execution_test.go @@ -10,7 +10,6 @@ import ( "time" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/google/uuid" ) @@ -24,10 +23,7 @@ func verifyNativePublicExecution(t *testing.T, h *dispatchHarness, parent contex } ctx, cancel := context.WithCancel(parent) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -40,7 +36,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, auth, "codex", workerExecution(worker), nativeDeploymentDefaults(h.s, "gpt-5.5", provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(worker), nativeDeploymentDefaults(h.s, "gpt-5.5", provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/path_id_semantics_public_test.go b/services/core/internal/store/path_id_semantics_public_test.go index bb671215..e831cedb 100644 --- a/services/core/internal/store/path_id_semantics_public_test.go +++ b/services/core/internal/store/path_id_semantics_public_test.go @@ -90,14 +90,14 @@ func TestMalformedPathIDsMatchMissingPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } @@ -359,7 +359,7 @@ func TestMalformedPathIDsMatchMissingPostgres(t *testing.T) { } // Storage availability checks also run before the lookup of a missing identifier. - h, err = publicHandler(t, store.New(pool), auth, "codex") + h, err = publicHandler(t, store.New(pool), fixtureDB{pool: pool}, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/prepared_dispatch_failure_test.go b/services/core/internal/store/prepared_dispatch_failure_test.go index 393f6660..7394d99c 100644 --- a/services/core/internal/store/prepared_dispatch_failure_test.go +++ b/services/core/internal/store/prepared_dispatch_failure_test.go @@ -9,6 +9,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) @@ -121,12 +122,14 @@ func TestPreparedDispatchHandlesStartRejectionAndPendingStartCancellation(t *tes } } -func TestPreparedDispatchRejectsPooledWriterBeforePreparation(t *testing.T) { +func TestPreparedDispatchRejectsClosedLeaseBeforePreparation(t *testing.T) { h, pending := preparedDispatchHarness(t) - h.d.Store = h.s - got, err := h.d.RunEnvironmentInput(context.Background(), h.tenant, h.session.ID, pending.ID) - if err == nil || got.Turn.ID != "" { - t.Fatal("pooled writer reached native preparation", got, err) + if err := h.lease.Close(t.Context()); err != nil { + t.Fatal(err) + } + got, err := h.d.RunEnvironmentInput(context.Background(), h.lease, h.tenant, h.session.ID, pending.ID) + if !errors.Is(err, pgunit.ErrLeaseClosed) || got.Turn.ID != "" { + t.Fatal("closed lease reached native preparation", got, err) } } diff --git a/services/core/internal/store/prepared_dispatch_test.go b/services/core/internal/store/prepared_dispatch_test.go index b9ab4573..8e4c7a4a 100644 --- a/services/core/internal/store/prepared_dispatch_test.go +++ b/services/core/internal/store/prepared_dispatch_test.go @@ -22,12 +22,8 @@ func preparedDispatchHarness(t *testing.T) (*dispatchHarness, store.EnvironmentI t.Helper() h := newDispatchHarnessForSession(t, []byte(`{"agent":{"model":"test-model","instructions":"Keep this instruction."},"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`), true) assertNoRuntimeAllocation(t, h) - writer, err := store.NewExecution(t.Context(), h.s) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = writer.CloseExecution(context.Background()) }) - h.d.Store = writer + owner := executionOwner(t, h.db, h.s) + h.d.Store, h.lease = owner.Store, owner.Lease enableWorkerEnvironment(t, h) pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []store.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}, {Kind: "message", Payload: json.RawMessage(`{"text":"second"}`)}}) if err != nil { @@ -39,7 +35,7 @@ func preparedDispatchHarness(t *testing.T) (*dispatchHarness, store.EnvironmentI func runPreparedDispatch(h *dispatchHarness, ctx context.Context, pending store.EnvironmentInputReservation) <-chan preparedDispatchResult { out := make(chan preparedDispatchResult, 1) go func() { - result, err := h.d.RunEnvironmentInput(ctx, h.tenant, h.session.ID, pending.ID) + result, err := h.d.RunEnvironmentInput(ctx, h.lease, h.tenant, h.session.ID, pending.ID) out <- preparedDispatchResult{result, err} }() return out @@ -114,7 +110,7 @@ func TestPreparedDispatchPromotesOriginalBatchAndPersistsCompletion(t *testing.T if err != nil || bound.NativeSessionID != "retained-prepared-native" { t.Fatal("native identity was not committed", bound, err) } - retry, err := h.d.RunEnvironmentInput(t.Context(), h.tenant, h.session.ID, pending.ID) + retry, err := h.d.RunEnvironmentInput(t.Context(), h.lease, h.tenant, h.session.ID, pending.ID) if err != nil || len(retry.Reservation.Receipts) != 2 || !retry.Reservation.Receipts[0].Replayed || retry.Reservation.Receipts[0].TurnID != start.RunID || retry.Turn.ID != "" { t.Fatal("replay executed again", retry, err) } diff --git a/services/core/internal/store/project_api_keys_http_test.go b/services/core/internal/store/project_api_keys_http_test.go index 0bd8d601..a64f773b 100644 --- a/services/core/internal/store/project_api_keys_http_test.go +++ b/services/core/internal/store/project_api_keys_http_test.go @@ -13,13 +13,13 @@ import ( ) func TestProjectAndSharedKeysHTTPManagement(t *testing.T) { - st, _ := store.NewTestStore(t) + st, db := newTestStoreDB(t) adminToken := uuid.NewString() admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential(adminToken)}) if err != nil { t.Fatal(err) } - h, err := publicHandler(t, st, nil, "codex", storeKeys(st), withCoreKeys(admin)) + h, err := publicHandler(t, st, db, nil, "codex", storeKeys(st), withCoreKeys(admin)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/public_execution_test.go b/services/core/internal/store/public_execution_test.go index 28c9a2ff..8433691b 100644 --- a/services/core/internal/store/public_execution_test.go +++ b/services/core/internal/store/public_execution_test.go @@ -8,7 +8,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -28,10 +27,7 @@ func TestExecutionWorkerAdmissionBindingAndRecovery(t *testing.T) { h.session = publicSession(t, h, "public") ctx, cancel := context.WithCancel(context.Background()) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() t.Cleanup(func() { @@ -42,7 +38,7 @@ func TestExecutionWorkerAdmissionBindingAndRecovery(t *testing.T) { t.Error("worker did not stop") } }) - if second, err := execution.StartWorker(ctx, h.d); err == nil { + if second, err := startWorkerErr(ctx, h.db, h.d); err == nil { cancel() go second.Run(ctx) t.Fatal("second service acquired database") @@ -151,10 +147,7 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { if _, err := h.s.SubmitMessage(ctx, h.tenant, queued.ID, "first", json.RawMessage(`{"text":"Not sent"}`)); err != nil { t.Fatal(err) } - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) stopped, cancel := context.WithCancel(ctx) cancel() if err := worker.Run(stopped); err != context.Canceled { @@ -175,7 +168,7 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { if err != nil || pending.LastTurn.Status != store.TurnCancelled { t.Fatal(pending, err) } - restarted, err := execution.StartWorker(ctx, h.d) + restarted, err := startWorkerErr(ctx, h.db, h.d) if err != nil { t.Fatal("lease not released", err) } diff --git a/services/core/internal/store/public_handler_fixture_test.go b/services/core/internal/store/public_handler_fixture_test.go index 9b03cc85..e196be00 100644 --- a/services/core/internal/store/public_handler_fixture_test.go +++ b/services/core/internal/store/public_handler_fixture_test.go @@ -22,11 +22,11 @@ import ( const testExecutorURL = "wss://core.example/api/v1/agent-daemon/ws" // publicHandler serves s through api.NewHandler. s backs every area the Store -// implements, 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.Store, keys fixtureKeyResolver, engine string, configure ...func(*api.Dependencies)) (http.Handler, error) { +// implements, and db is the database and credential key that built s. 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.Store, db fixtureDB, keys fixtureKeyResolver, engine string, configure ...func(*api.Dependencies)) (http.Handler, error) { t.Helper() admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential("admin")}) if err != nil { diff --git a/services/core/internal/store/remote_mcp_credentials_test.go b/services/core/internal/store/remote_mcp_credentials_test.go index 8f3bf11e..60eea2c4 100644 --- a/services/core/internal/store/remote_mcp_credentials_test.go +++ b/services/core/internal/store/remote_mcp_credentials_test.go @@ -13,20 +13,20 @@ import ( func TestSelfHostedServiceMCPRejectionDoesNotRequireCredentialDecryption(t *testing.T) { for _, mode := range []string{"missing key", "deleted", "tampered"} { t.Run(mode, func(t *testing.T) { - s, pool, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) + s, db, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) switch mode { case "missing key": - s = store.New(pool) + s, db = store.New(db.pool), fixtureDB{pool: db.pool} case "deleted": if _, err := s.DeleteCredential(t.Context(), tenant, vault.ID, credential.ID); err != nil { t.Fatal(err) } case "tampered": - if _, err := 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 := 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 { t.Fatal(err) } } - handler := selfHostedMCPAdmissionHandler(t, s, tenant) + handler := selfHostedMCPAdmissionHandler(t, s, db, tenant) for _, initial := range []bool{false, true} { body := map[string]any{ "agent": map[string]any{"model": "model", "tools": []any{map[string]any{ @@ -60,7 +60,7 @@ func TestSelfHostedServiceMCPRejectionDoesNotRequireCredentialDecryption(t *test if !strings.Contains(response.Body.String(), message) { t.Fatal("unsupported placement attempted credential decryption", response.Body) } - assertSelfHostedMCPRejectionHasNoWrites(t, pool, tenant) + assertSelfHostedMCPRejectionHasNoWrites(t, db.pool, tenant) } }) } diff --git a/services/core/internal/store/remote_mcp_test.go b/services/core/internal/store/remote_mcp_test.go index 0f543ffd..0c324b1e 100644 --- a/services/core/internal/store/remote_mcp_test.go +++ b/services/core/internal/store/remote_mcp_test.go @@ -15,8 +15,8 @@ import ( ) func TestSelfHostedServiceMCPRejectedWithoutWrites(t *testing.T) { - s, pool, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) - handler := selfHostedMCPAdmissionHandler(t, s, tenant) + s, db, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) + handler := selfHostedMCPAdmissionHandler(t, s, db, 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,19 +67,19 @@ func TestSelfHostedServiceMCPRejectedWithoutWrites(t *testing.T) { t.Fatal("rejected request exposed private authentication") } - assertSelfHostedMCPRejectionHasNoWrites(t, pool, tenant) + assertSelfHostedMCPRejectionHasNoWrites(t, db.pool, tenant) } } } -func selfHostedMCPAdmissionFixture(t *testing.T) (*store.Store, *pgxpool.Pool, string, store.Vault, store.Credential) { +func selfHostedMCPAdmissionFixture(t *testing.T) (*store.Store, fixtureDB, string, store.Vault, store.Credential) { t.Helper() _, pool := store.NewTestStore(t) cipher, err := credentialcrypto.New([]byte(strings.Repeat("k", 32))) if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} tenant := uuid.NewString() vault, err := s.CreateVault(t.Context(), tenant, store.CreateVaultInput{}) if err != nil { @@ -89,13 +89,13 @@ func selfHostedMCPAdmissionFixture(t *testing.T) (*store.Store, *pgxpool.Pool, s if err != nil { t.Fatal(err) } - return s, pool, tenant, vault, credential + return s, db, tenant, vault, credential } -func selfHostedMCPAdmissionHandler(t *testing.T, s *store.Store, tenant string) http.Handler { +func selfHostedMCPAdmissionHandler(t *testing.T, s *store.Store, db fixtureDB, 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, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) + handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/request_body_public_test.go b/services/core/internal/store/request_body_public_test.go index 6ce2663f..74baf2b2 100644 --- a/services/core/internal/store/request_body_public_test.go +++ b/services/core/internal/store/request_body_public_test.go @@ -29,7 +29,7 @@ func TestRequestBodyGateRejectsWithoutWritesPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, s), unavailable, acceptUnavailable(t)) + h, err := publicHandler(t, s, db, 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 := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/runtime_allocations.go b/services/core/internal/store/runtime_allocations.go index 0c433b19..8e7cb304 100644 --- a/services/core/internal/store/runtime_allocations.go +++ b/services/core/internal/store/runtime_allocations.go @@ -148,7 +148,7 @@ func (s *Store) GetRuntimeAllocation(ctx context.Context, tenant, environment st // ListRuntimeAllocations retains unresolved cleanup in bounded recovery scans. func (s *Store) ListRuntimeAllocations(ctx context.Context, after string) ([]RuntimeAllocation, error) { - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return nil, err } id := pgtype.UUID{Valid: true} @@ -228,7 +228,7 @@ type UnallocatedHostedEnvironment struct { } func (s *Store) ListUnallocatedHostedEnvironments(ctx context.Context, after string) ([]UnallocatedHostedEnvironment, error) { - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return nil, err } id := pgtype.UUID{Valid: true} diff --git a/services/core/internal/store/runtime_allocations_test.go b/services/core/internal/store/runtime_allocations_test.go index 575e72bc..f25ac548 100644 --- a/services/core/internal/store/runtime_allocations_test.go +++ b/services/core/internal/store/runtime_allocations_test.go @@ -31,7 +31,7 @@ func TestRuntimeAllocationAtomicOwnershipAndRecovery(t *testing.T) { if _, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, uuid.NewString(), runtimedevice.HashCredential(secret)); !errors.Is(err, ErrIdempotencyConflict) { t.Fatalf("provider target changed: %v", err) } - if err := w.CloseExecution(t.Context()); err != nil { + if err := w.lease.Close(t.Context()); err != nil { t.Fatal(err) } if _, err := w.ObserveRuntimeRunning(t.Context(), owner); err == nil { diff --git a/services/core/internal/store/runtime_capabilities_pending_test.go b/services/core/internal/store/runtime_capabilities_pending_test.go index 5c687ed4..0c6902cc 100644 --- a/services/core/internal/store/runtime_capabilities_pending_test.go +++ b/services/core/internal/store/runtime_capabilities_pending_test.go @@ -19,7 +19,7 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} tenant := uuid.NewString() session, err := s.CreateSession(t.Context(), tenant, store.CreateSessionInput{ Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), @@ -35,7 +35,7 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { } provider := &initializingProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, initializationPeer: initializationPeer{deferred: true}} key := uuid.NewString() - worker, _ := managedWorkerMode(t, s, key, provider, false, true) + worker, _ := managedWorkerMode(t, s, db, key, provider, false, true) owner, err := worker.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil || initializationState(t, s, owner.TenantID, owner.EnvironmentID) != "pending" { t.Fatal(owner, err) diff --git a/services/core/internal/store/runtime_compute_lifecycle_test.go b/services/core/internal/store/runtime_compute_lifecycle_test.go index 0626d5f1..ea36ef8f 100644 --- a/services/core/internal/store/runtime_compute_lifecycle_test.go +++ b/services/core/internal/store/runtime_compute_lifecycle_test.go @@ -21,7 +21,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" "github.com/gorilla/websocket" - "github.com/jackc/pgx/v5/pgxpool" ) // The controlled provider records external effects independently of DB phases. @@ -259,7 +258,7 @@ func (p *fakeCheckpointProvider) connect(ctx context.Context, b sandbox.Bootstra type computeLifecycleFixture struct { t *testing.T store *store.Store - pool *pgxpool.Pool + db fixtureDB provider *fakeCheckpointProvider worker *execution.Worker stop func() @@ -269,7 +268,7 @@ type computeLifecycleFixture struct { func newComputeLifecycleFixture(t *testing.T, maxActive, maxRetained int) *computeLifecycleFixture { t.Helper() - s, pool := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(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(s), Registry: registry}) @@ -283,17 +282,14 @@ func newComputeLifecycleFixture(t *testing.T, maxActive, maxRetained int) *compu } server.Close() }) - f := &computeLifecycleFixture{t: t, store: s, pool: pool, provider: p, key: uuid.NewString(), policy: execution.RuntimeSuspensionPolicy{IdleTimeout: time.Second, Retention: time.Hour, MaxActive: maxActive, MaxRetained: maxRetained}} + 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.start() return f } func (f *computeLifecycleFixture) start() { t := f.t t.Helper() - w, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: f.store, Registry: f.provider.registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: f.key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: f.provider, Suspension: &f.policy}}) - if err != nil { - t.Fatal(err) - } + w := startWorker(t, t.Context(), f.db, &execution.Dispatcher{Store: f.store, Registry: f.provider.registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: f.key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: f.provider, Suspension: &f.policy}}) var once sync.Once stop := func() { once.Do(func() { ctx, cancel := context.WithCancel(context.Background()); cancel(); _ = w.Run(ctx) }) @@ -303,7 +299,7 @@ func (f *computeLifecycleFixture) start() { } func (f *computeLifecycleFixture) sql(query string, args ...any) { f.t.Helper() - if _, err := f.pool.Exec(f.t.Context(), query, args...); err != nil { + if _, err := f.db.pool.Exec(f.t.Context(), query, args...); err != nil { f.t.Fatal(err) } } @@ -377,7 +373,7 @@ func TestRuntimeComputeLifecycleIdleSuspendAndQueuedSameSessionWake(t *testing.T t.Fatal("wake replaced Session or replayed allocation") } var completedCount, queuedCount int - if err := f.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.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 { t.Fatal("wake replayed/consumed prior or next Turn", err) } if got, err := f.store.GetSession(t.Context(), tenant, session.ID); err != nil || string(got.Configuration) != string(session.Configuration) { diff --git a/services/core/internal/store/runtime_configuration_cleanup_test.go b/services/core/internal/store/runtime_configuration_cleanup_test.go index 32eddbe1..242d98af 100644 --- a/services/core/internal/store/runtime_configuration_cleanup_test.go +++ b/services/core/internal/store/runtime_configuration_cleanup_test.go @@ -5,7 +5,6 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -82,7 +81,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, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) key := uuid.NewString() p := &configurationCleanupProvider{ lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}, loseCreate: test.loseCreate}, @@ -90,7 +89,7 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { settleInspection: test.settleInspection, foreign: test.foreign, inspectionError: test.inspectionError, killError: test.killError, } - w, _ := managedWorker(t, s, key, p) + w, _ := managedWorker(t, s, db, 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 == "" { diff --git a/services/core/internal/store/runtime_connection_test.go b/services/core/internal/store/runtime_connection_test.go index 1e2e5992..5c589955 100644 --- a/services/core/internal/store/runtime_connection_test.go +++ b/services/core/internal/store/runtime_connection_test.go @@ -20,7 +20,7 @@ import ( ) func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) tenant, session, environment := managedSession(t, s) server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" @@ -34,10 +34,7 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { p := &lifecycleProvider{resources: map[string]sandbox.Info{}} key := uuid.NewString() start := func() *execution.Worker { - w, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: server.URL + "/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p}}) - if err != nil { - t.Fatal(err) - } + w := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: server.URL + "/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p}}) return w } stop := func(w *execution.Worker) { diff --git a/services/core/internal/store/runtime_creation_settlement_test.go b/services/core/internal/store/runtime_creation_settlement_test.go index 62aade0c..6e1460d9 100644 --- a/services/core/internal/store/runtime_creation_settlement_test.go +++ b/services/core/internal/store/runtime_creation_settlement_test.go @@ -6,7 +6,6 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -40,10 +39,10 @@ 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, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) key := uuid.NewString() p := &absentCreationProvider{} - w, _ := managedWorker(t, s, key, p) + w, _ := managedWorker(t, s, db, key, p) tenant, session, environment := managedSession(t, s) ctx, cancel := context.WithCancel(t.Context()) defer cancel() @@ -78,10 +77,10 @@ func TestManagedRuntimeConfirmedAbsentCreateReleasesAtomically(t *testing.T) { } } func TestManagedRuntimeForeignAbsenceCannotReleaseCreation(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) key := uuid.NewString() p := &absentCreationProvider{foreign: true} - w, _ := managedWorker(t, s, key, p) + w, _ := managedWorker(t, s, db, key, p) tenant, _, environment := managedSession(t, s) if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err == nil { t.Fatal("foreign absence accepted") @@ -92,10 +91,10 @@ func TestManagedRuntimeForeignAbsenceCannotReleaseCreation(t *testing.T) { } } func TestManagedRuntimeObservedSettlementAllowsOwnedCleanup(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) key := uuid.NewString() p := &absentCreationProvider{observeSettled: true} - w, _ := managedWorker(t, s, key, p) + w, _ := managedWorker(t, s, db, 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") diff --git a/services/core/internal/store/runtime_deployment_worker_test.go b/services/core/internal/store/runtime_deployment_worker_test.go index 3df49f0a..3570c6cb 100644 --- a/services/core/internal/store/runtime_deployment_worker_test.go +++ b/services/core/internal/store/runtime_deployment_worker_test.go @@ -7,15 +7,14 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) func TestManagedDeploymentStartupRejectsSwitchBeforeBackendAccess(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) key := uuid.NewString() old := &lifecycleProvider{resources: map[string]sandbox.Info{}} - worker, stop := managedWorker(t, s, key, old) + worker, stop := managedWorker(t, s, db, key, old) tenant, _, environment := managedSession(t, s) owner, err := worker.ProvisionEnvironment(t.Context(), tenant, environment.ID, key) if err != nil { @@ -25,13 +24,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 := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: config}) + _, err := startWorkerErr(t.Context(), db, &execution.Dispatcher{Store: s, 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, key, old, true) + worker, stop = managedWorkerMode(t, s, db, 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) @@ -51,6 +50,6 @@ func TestManagedDeploymentStartupRejectsSwitchBeforeBackendAccess(t *testing.T) t.Fatal("rejected startup rewrote resource owner", got, err) } // Failed startup relinquishes its lease, so the original backend can resume. - _, stop = managedWorker(t, s, key, old) + _, stop = managedWorker(t, s, db, key, old) stop() } diff --git a/services/core/internal/store/runtime_enrollment_connection_test.go b/services/core/internal/store/runtime_enrollment_connection_test.go index 0be93bad..0f880671 100644 --- a/services/core/internal/store/runtime_enrollment_connection_test.go +++ b/services/core/internal/store/runtime_enrollment_connection_test.go @@ -21,7 +21,7 @@ import ( ) func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { - s, pool := store.NewTestStore(t) + s, db := newTestStoreDB(t) principal := store.FixtureExecutorPrincipal(t, s, uuid.NewString()) session, err := s.CreateSession(t.Context(), principal.TenantID, store.CreateSessionInput{ Creator: principal.Subject(), Engine: "codex", IdempotencyKey: uuid.NewString(), @@ -83,10 +83,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { server.Start() t.Cleanup(func() { server.Close(); runtime.CloseConnections(registry) }) start := func() func() { - worker, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: registry}) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: registry}) ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -148,7 +145,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { second := connect(rotated.Token) await("connected") assertConnection(environment.ID, rotated.Token, "connected", 200) - awaitRelease := observeExecutionLeaseRelease(t, pool) + awaitRelease := observeExecutionLeaseRelease(t, db.pool) stop() stop = nil awaitRelease() @@ -166,7 +163,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { t.Fatal("revoked socket retained authority") } var allocations int - if err = pool.QueryRow(t.Context(), "SELECT count(*) FROM runtime_allocations WHERE environment_id=$1", environment.ID).Scan(&allocations); err != nil || allocations != 0 { + if err = db.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 := s.GetSession(t.Context(), principal.TenantID, session.ID) diff --git a/services/core/internal/store/runtime_initialization_test.go b/services/core/internal/store/runtime_initialization_test.go index bd35b2c2..8b445181 100644 --- a/services/core/internal/store/runtime_initialization_test.go +++ b/services/core/internal/store/runtime_initialization_test.go @@ -49,7 +49,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} tenant := uuid.NewString() input := store.CreateSessionInput{Creator: store.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 { @@ -84,7 +84,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { return completedInitialization(proto.RuntimePreparePayload{}, nil) } key := uuid.NewString() - w, stop := managedWorkerMode(t, s, key, p, false, true) + w, stop := managedWorkerMode(t, s, db, key, p, false, true) if mode == "restart" { awaitInitialization(t, s, tenant, env.ID, "failed") if p.writes.Load() != 0 { @@ -120,7 +120,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { t.Fatal("completed preparation blocked", err) } stop() - _, _ = managedWorkerMode(t, s, key, p, false, true) + _, _ = managedWorkerMode(t, s, db, key, p, false, true) time.Sleep(350 * time.Millisecond) if int(p.writes.Load()) != expectedSteps { t.Fatal("completed preparation replayed") @@ -139,7 +139,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { } func TestManagedRuntimePreparationAllOperationsUsePeer(t *testing.T) { - s := hostedFailureStore(t) + s, db := 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."} { @@ -180,7 +180,7 @@ func TestManagedRuntimePreparationAllOperationsUsePeer(t *testing.T) { return completedInitialization(request, data) } key := uuid.NewString() - worker, _ := managedWorkerMode(t, s, key, provider, false, true) + worker, _ := managedWorkerMode(t, s, db, key, provider, false, true) if _, err := worker.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/runtime_input_admission_test.go b/services/core/internal/store/runtime_input_admission_test.go index b616529b..972450d4 100644 --- a/services/core/internal/store/runtime_input_admission_test.go +++ b/services/core/internal/store/runtime_input_admission_test.go @@ -12,7 +12,7 @@ import ( ) func TestManagedRuntimeMaintenancePreservesCancelAndRetry(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) tenant, session, _ := managedSession(t, s) inputs := []store.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"accepted work"}`)}} accepted, err := s.SubmitInputs(t.Context(), tenant, session.ID, "work", inputs) @@ -20,7 +20,7 @@ func TestManagedRuntimeMaintenancePreservesCancelAndRetry(t *testing.T) { t.Fatal(err) } p := &lifecycleProvider{resources: map[string]sandbox.Info{}} - w, stop := managedWorkerMode(t, s, uuid.NewString(), p, true) + w, stop := managedWorkerMode(t, s, db, uuid.NewString(), p, true) defer stop() if _, err := w.CreateSession(t.Context(), tenant, store.CreateSessionInput{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: session.Configuration}); !errors.Is(err, store.ErrEnvironmentUnavailable) { t.Fatal("maintenance accepted new hosted Session", err) diff --git a/services/core/internal/store/runtime_lifecycle_nodes.go b/services/core/internal/store/runtime_lifecycle_nodes.go index 90b7fb4b..bfc3cc01 100644 --- a/services/core/internal/store/runtime_lifecycle_nodes.go +++ b/services/core/internal/store/runtime_lifecycle_nodes.go @@ -15,7 +15,7 @@ import ( func (s *Store) ListRuntimeLifecycleNodes(ctx context.Context) ([]string, error) { ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) defer cancel() - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return nil, err } rows, err := s.queries.ListRuntimeLifecycleNodes(ctx) @@ -53,7 +53,7 @@ func (s *Store) ListRuntimeAllocationsForNode(ctx context.Context, node, after s } ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) defer cancel() - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return nil, err } rows, err := s.queries.ListRuntimeAllocationsForNode(ctx, sqlc.ListRuntimeAllocationsForNodeParams{NodeID: nodeID, AfterID: afterID}) @@ -76,7 +76,7 @@ func (s *Store) ListUnallocatedHostedEnvironmentsForNode(ctx context.Context, no } ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) defer cancel() - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return nil, err } rows, err := s.queries.ListUnallocatedHostedEnvironmentsForNode(ctx, sqlc.ListUnallocatedHostedEnvironmentsForNodeParams{NodeID: nodeID, AfterID: afterID}) @@ -99,7 +99,7 @@ func (s *Store) ResolveRuntimeLifecycleNode(ctx context.Context, tenant, environ } ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) defer cancel() - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return "", err } row, err := s.queries.GetRuntimeLifecyclePlacement(ctx, sqlc.GetRuntimeLifecyclePlacementParams{TenantID: lookup.TenantID, ID: lookup.ID}) diff --git a/services/core/internal/store/runtime_lifecycle_nodes_test.go b/services/core/internal/store/runtime_lifecycle_nodes_test.go index 8d05b176..7c4e6235 100644 --- a/services/core/internal/store/runtime_lifecycle_nodes_test.go +++ b/services/core/internal/store/runtime_lifecycle_nodes_test.go @@ -230,7 +230,7 @@ func TestRuntimeLifecycleLegacyLaneAndOwnerLoss(t *testing.T) { if _, err := w.ListUnallocatedHostedEnvironmentsForNode(t.Context(), "", "bad"); err == nil { t.Fatal("invalid cursor accepted") } - if err := w.CloseExecution(t.Context()); err != nil { + if err := w.lease.Close(t.Context()); err != nil { t.Fatal(err) } if _, err := w.ListRuntimeLifecycleNodes(t.Context()); err == nil { diff --git a/services/core/internal/store/runtime_lifecycle_test.go b/services/core/internal/store/runtime_lifecycle_test.go index d976cb02..42d0abdf 100644 --- a/services/core/internal/store/runtime_lifecycle_test.go +++ b/services/core/internal/store/runtime_lifecycle_test.go @@ -71,12 +71,12 @@ func (p *lifecycleProvider) RunCommand(context.Context, sandbox.Reference, sandb return sandbox.CommandResult{}, errors.New("not used") } -func managedWorker(t *testing.T, s *store.Store, key string, p sandbox.SandboxProvider) (*execution.Worker, func()) { +func managedWorker(t *testing.T, s *store.Store, db fixtureDB, key string, p sandbox.SandboxProvider) (*execution.Worker, func()) { t.Helper() - return managedWorkerMode(t, s, key, p, false) + return managedWorkerMode(t, s, db, key, p, false) } -func managedWorkerMode(t *testing.T, s *store.Store, key string, p sandbox.SandboxProvider, maintenance bool, run ...bool) (*execution.Worker, func()) { +func managedWorkerMode(t *testing.T, s *store.Store, db fixtureDB, key string, p sandbox.SandboxProvider, maintenance bool, run ...bool) (*execution.Worker, func()) { t.Helper() registry := runtimegateway.NewRegistry() if peer, ok := p.(interface { @@ -87,10 +87,7 @@ func managedWorkerMode(t *testing.T, s *store.Store, key string, p sandbox.Sandb t.Cleanup(server.Close) peer.setRuntimeGateway(t, "ws"+strings.TrimPrefix(server.URL, "http"), registry) } - w, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: registry, ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p, AdmissionPaused: maintenance}}) - if err != nil { - t.Fatal(err) - } + w := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, 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) @@ -147,11 +144,11 @@ func reconcileManagedState(t *testing.T, w *execution.Worker, s *store.Store, te } func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) tenant, session, env := managedSession(t, s) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}, loseCreate: true} - w, stop := managedWorker(t, s, key, p) + w, stop := managedWorker(t, s, db, key, p) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err == nil || owner.ID == "" { t.Fatal("fault did not retain allocation") @@ -161,7 +158,7 @@ func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { t.Fatal("provider received unbound credential") } stop() - next, _ := managedWorker(t, s, key, p) + next, _ := managedWorker(t, s, db, key, p) reconcileManagedState(t, next, s, tenant, env.ID, "running") recovered, err := s.GetRuntimeAllocation(t.Context(), tenant, env.ID) if err != nil || recovered.ID != owner.ID || !recovered.CreateSettled || recovered.State != "running" { @@ -186,11 +183,11 @@ func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { } func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(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, key, p) + w, _ := managedWorker(t, s, db, key, p) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err == nil { t.Fatal("expected uncertain creation") @@ -216,16 +213,16 @@ func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { } func TestManagedRuntimeExpiryRevokesWhenProviderUnavailable(t *testing.T) { - s, pool := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) tenant, _, env := managedSession(t, s) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} - w, _ := managedWorker(t, s, key, p) + w, _ := managedWorker(t, s, db, key, p) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) } - if _, err := pool.Exec(t.Context(), "UPDATE runtime_allocations SET kept_at=clock_timestamp()-interval '61 minutes' WHERE id=$1", owner.ID); err != nil { + 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 { t.Fatal(err) } p.unavailable = true @@ -243,11 +240,11 @@ func TestManagedRuntimeExpiryRevokesWhenProviderUnavailable(t *testing.T) { } func TestManagedRuntimeStoppedComputeDoesNotRequestCleanup(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) tenant, _, env := managedSession(t, s) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} - w, _ := managedWorker(t, s, key, p) + w, _ := managedWorker(t, s, db, key, p) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) diff --git a/services/core/internal/store/runtime_node_lifecycle_fixture_test.go b/services/core/internal/store/runtime_node_lifecycle_fixture_test.go index 9dbbdba5..77eaa43a 100644 --- a/services/core/internal/store/runtime_node_lifecycle_fixture_test.go +++ b/services/core/internal/store/runtime_node_lifecycle_fixture_test.go @@ -94,7 +94,7 @@ func newNodeIsolationFixture(t *testing.T, mode string) *nodeIsolationFixture { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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{})} @@ -126,10 +126,7 @@ func newNodeIsolationFixture(t *testing.T, mode string) *nodeIsolationFixture { // 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, err = execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, 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}}) - if err != nil { - t.Fatal(err) - } + f.worker = startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, 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 := store.SandboxDeploymentTestSpec("microsandbox") raw, _ := json.Marshal(spec) if _, err := pool.Exec(t.Context(), "UPDATE runtime_deployment SET specification=$1", raw); err != nil { diff --git a/services/core/internal/store/runtime_pending_test.go b/services/core/internal/store/runtime_pending_test.go index 26bc8fed..96fc3f6d 100644 --- a/services/core/internal/store/runtime_pending_test.go +++ b/services/core/internal/store/runtime_pending_test.go @@ -13,7 +13,7 @@ import ( ) func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) tenant, idle, idleEnvironment := managedSession(t, s) initial, err := s.CreateSession(t.Context(), tenant, store.CreateSessionInput{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"model":"test"},"environment":{"type":"openai_hosted"}}`), InitialInputs: []store.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"hello"}`)}}}) if err != nil { @@ -26,10 +26,7 @@ func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} start := func() *execution.Worker { - w, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: &execution.RuntimeProvider{CoreURL: "http://core.invalid/api/v1", InstallationID: key, BackendFingerprint: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", Provider: p}}) - if err != nil { - t.Fatal(err) - } + w := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, 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) { diff --git a/services/core/internal/store/runtime_scan_test.go b/services/core/internal/store/runtime_scan_test.go index 5adb0322..7cd13fd6 100644 --- a/services/core/internal/store/runtime_scan_test.go +++ b/services/core/internal/store/runtime_scan_test.go @@ -11,7 +11,6 @@ import ( "github.com/google/uuid" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) type scanProvider struct { @@ -27,10 +26,10 @@ 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, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) p := &scanProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}} key := uuid.NewString() - w, _ := managedWorker(t, s, key, p) + w, _ := managedWorker(t, s, db, key, p) var ids []string for range count { tenant, _, env := managedSession(t, s) @@ -66,10 +65,10 @@ func TestManagedRuntimeScanWrapServicesNextPage(t *testing.T) { } func TestManagedRuntimeScanEmptyAfterCleanupAndCanceledCall(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) p := &scanProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}} key := uuid.NewString() - w, _ := managedWorker(t, s, key, p) + w, _ := managedWorker(t, s, db, key, p) tenant, session, env := managedSession(t, s) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { diff --git a/services/core/internal/store/runtime_suspension.go b/services/core/internal/store/runtime_suspension.go index 4661292e..2ff0a10c 100644 --- a/services/core/internal/store/runtime_suspension.go +++ b/services/core/internal/store/runtime_suspension.go @@ -99,7 +99,7 @@ func (s *Store) RuntimeActivity(ctx context.Context, owner RuntimeAllocation) (R if err != nil { return RuntimeActivity{}, err } - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return RuntimeActivity{}, err } row, err := s.queries.GetRuntimeActivity(ctx, id) @@ -143,7 +143,7 @@ func (s *Store) ClearRuntimeWake(ctx context.Context, owner RuntimeAllocation, o if err != nil { return err } - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return err } return s.queries.ClearRuntimeWake(ctx, sqlc.ClearRuntimeWakeParams{ID: id, ComputeActivityAt: pgtype.Timestamptz{Time: observedActivity, Valid: true}}) @@ -154,7 +154,7 @@ func (s *Store) CountRuntimeComputeReservations(ctx context.Context, provider st if err != nil { return 0, err } - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return 0, err } return s.queries.CountRuntimeComputeReservations(ctx, id) @@ -165,7 +165,7 @@ func (s *Store) CountRuntimeRetainedAllocations(ctx context.Context, provider st if err != nil { return 0, err } - if err := s.CheckExecutionOwnership(ctx); err != nil { + if err := s.checkExecutionOwnership(ctx); err != nil { return 0, err } return s.queries.CountRuntimeRetainedAllocations(ctx, id) diff --git a/services/core/internal/store/runtime_suspension_test.go b/services/core/internal/store/runtime_suspension_test.go index 38b7b876..03abe136 100644 --- a/services/core/internal/store/runtime_suspension_test.go +++ b/services/core/internal/store/runtime_suspension_test.go @@ -359,7 +359,7 @@ func TestRuntimeSuspensionExpiredRunningAndLostWriterAreFenced(t *testing.T) { if _, err := w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond); !errors.Is(err, ErrTurnConflict) { t.Fatal("expired running allocation entered checkpoint", err) } - if err := w.CloseExecution(t.Context()); err != nil { + if err := w.lease.Close(t.Context()); err != nil { t.Fatal(err) } if _, err := w.RuntimeActivity(t.Context(), owner); err == nil { diff --git a/services/core/internal/store/runtime_wake_hint_integration_test.go b/services/core/internal/store/runtime_wake_hint_integration_test.go index e2c4c069..5240d78e 100644 --- a/services/core/internal/store/runtime_wake_hint_integration_test.go +++ b/services/core/internal/store/runtime_wake_hint_integration_test.go @@ -72,7 +72,7 @@ func newWakeHintIntegration(t *testing.T) *wakeHintIntegration { fakeCheckpointProvider: f.provider, sentinel: sentinel.owner.ID, release: make(chan struct{}), scans: make(chan int, 16), } - worker, err := execution.StartWorker(t.Context(), &execution.Dispatcher{ + worker := startWorker(t, t.Context(), f.db, &execution.Dispatcher{ Store: f.store, Registry: f.provider.registry, ManagedRuntimes: &execution.RuntimeProvider{ CoreURL: "http://core.invalid/api/v1", InstallationID: f.key, @@ -80,9 +80,6 @@ func newWakeHintIntegration(t *testing.T) *wakeHintIntegration { Provider: provider, Suspension: &f.policy, }, }) - if err != nil { - t.Fatal(err) - } ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) var once sync.Once @@ -123,7 +120,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.pool.QueryRow(t.Context(), + return f.fixture.db.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 }) @@ -181,7 +178,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.pool.QueryRow(t.Context(), + if err := f.fixture.db.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/internal/store/runtime_worker_recovery_test.go b/services/core/internal/store/runtime_worker_recovery_test.go index 5d257b61..417227c6 100644 --- a/services/core/internal/store/runtime_worker_recovery_test.go +++ b/services/core/internal/store/runtime_worker_recovery_test.go @@ -34,12 +34,8 @@ func runtimeWorkerHarness(t *testing.T) (*dispatchHarness, *pgxpool.Pool) { func TestPreparedDispatchKeepsPendingReservationAfterComputeConflict(t *testing.T) { h, _ := runtimeWorkerHarness(t) - writer, err := store.NewExecution(t.Context(), h.s) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = writer.CloseExecution(context.Background()) }) - h.d.Store = writer + owner := executionOwner(t, h.db, h.s) + h.d.Store, h.lease = owner.Store, owner.Lease pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []store.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) if err != nil { t.Fatal(err) @@ -65,10 +61,7 @@ func TestWorkerWaitsForComputeAndSurvivesPromotionConflict(t *testing.T) { if err != nil { t.Fatal(err) } - worker, err := execution.StartWorker(t.Context(), h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), h.db, h.d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -152,7 +145,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 := execution.StartWorker(t.Context(), h.d) + worker, err := startWorkerErr(t.Context(), h.db, h.d) if err != nil { t.Fatal("queued wake blocked Core startup", err) } diff --git a/services/core/internal/store/sandbox_deployment_setup_test.go b/services/core/internal/store/sandbox_deployment_setup_test.go index a1b67d01..15620df8 100644 --- a/services/core/internal/store/sandbox_deployment_setup_test.go +++ b/services/core/internal/store/sandbox_deployment_setup_test.go @@ -70,7 +70,7 @@ func TestSandboxDeploymentSetupPersistsWithoutExecution(t *testing.T) { if err := pool.QueryRow(t.Context(), "SELECT (SELECT count(*) FROM sessions)+(SELECT count(*) FROM runtime_nodes)+(SELECT count(*) FROM runtime_allocations)+(SELECT count(*) FROM runtime_placements)").Scan(&sideEffects); err != nil || sideEffects != 0 { t.Fatal("setup or rejected admission created execution state", sideEffects, err) } - if err := w.CloseExecution(context.Background()); err != nil { + if err := w.lease.Close(context.Background()); err != nil { t.Fatal(err) } restarted := executionWriter(t, s) diff --git a/services/core/internal/store/sandbox_deployment_switch_worker_test.go b/services/core/internal/store/sandbox_deployment_switch_worker_test.go index b60453c6..7d422407 100644 --- a/services/core/internal/store/sandbox_deployment_switch_worker_test.go +++ b/services/core/internal/store/sandbox_deployment_switch_worker_test.go @@ -23,7 +23,7 @@ import ( func TestSandboxWorkerSwitchesAndRecoversFailedActivation(t *testing.T) { _, pool := store.NewManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{7}, 32)) - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} id := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} var fail atomic.Bool @@ -41,10 +41,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, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) - if err != nil { - t.Fatal(err) - } + w := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- w.Run(ctx) }() diff --git a/services/core/internal/store/sandbox_deployment_worker_test.go b/services/core/internal/store/sandbox_deployment_worker_test.go index 37ac7d79..a8bd14f2 100644 --- a/services/core/internal/store/sandbox_deployment_worker_test.go +++ b/services/core/internal/store/sandbox_deployment_worker_test.go @@ -16,7 +16,7 @@ import ( ) func TestSandboxDeploymentWorkerActivatesWithoutRestart(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) id := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} configuration := execution.NewDeferredRuntimeProvider(id, func(ctx context.Context) (*execution.RuntimeProvider, error) { @@ -31,10 +31,7 @@ func TestSandboxDeploymentWorkerActivatesWithoutRestart(t *testing.T) { }) start := func() (*execution.Worker, func()) { t.Helper() - w, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) - if err != nil { - t.Fatal(err) - } + w := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) var once sync.Once stop := func() { once.Do(func() { ctx, cancel := context.WithCancel(context.Background()); cancel(); _ = w.Run(ctx) }) diff --git a/services/core/internal/store/sandbox_node_auth_order_http_test.go b/services/core/internal/store/sandbox_node_auth_order_http_test.go index 85619cce..8bc3ca6d 100644 --- a/services/core/internal/store/sandbox_node_auth_order_http_test.go +++ b/services/core/internal/store/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, pool := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential(uuid.NewString())}) if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, s, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, s), withCoreKeys(admin)) + handler, err := publicHandler(t, s, db, 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 := 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 := 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 { 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 := pool.Exec(t.Context(), "UPDATE runtime_deployment SET installation_id=$1, web_managed=true WHERE singleton=true", claimed); err != nil { + if _, err := db.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/internal/store/saved_web_search_public_test.go b/services/core/internal/store/saved_web_search_public_test.go index 368cae54..e560ec78 100644 --- a/services/core/internal/store/saved_web_search_public_test.go +++ b/services/core/internal/store/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, pool := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } @@ -171,7 +171,7 @@ func TestSavedWebSearchPostgres(t *testing.T) { enabledTools := map[string]string{"mode-live": agents["mode-live"], "mode-cached": agents["mode-cached"], "type-only": agents["type-only"], "updated-to-live": agents["mode-disabled"]} // W4: every creation mode rejects enabled saved search without writes. - before := databaseDigest(t, pool) + before := databaseDigest(t, db.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{ @@ -223,7 +223,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, pool); !mapsEqual(before, after) { + if after := databaseDigest(t, db.pool); !mapsEqual(before, after) { t.Fatal("rejected Session creation changed persisted state") } diff --git a/services/core/internal/store/self_hosted_cancel_public_test.go b/services/core/internal/store/self_hosted_cancel_public_test.go index 46882ad5..911dba4d 100644 --- a/services/core/internal/store/self_hosted_cancel_public_test.go +++ b/services/core/internal/store/self_hosted_cancel_public_test.go @@ -21,7 +21,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, pool := store.NewModelTestStore(t) + s, db := newModelTestStoreDB(t) tenant, foreignTenant := uuid.NewString(), uuid.NewString() token, foreign := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ @@ -30,8 +30,8 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { }) serve := func() (*httptest.Server, func(bool)) { t.Helper() - worker, stop := publicInitialWorker(t, s) - handler, err := publicHandler(t, s, auth, "codex", workerExecution(worker), executorURL("https://offline-executor.example")) + worker, stop := publicInitialWorker(t, s, db) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(worker), executorURL("https://offline-executor.example")) if err != nil { t.Fatal(err) } @@ -75,7 +75,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { settings["accepted"] = accepted receipts := func(key, target string) []store.InputReceipt { t.Helper() - rows, err := pool.Query(t.Context(), `SELECT sequence, COALESCE(turn_id::text,'') FROM turn_inputs + rows, err := db.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) @@ -100,7 +100,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { snapshot := func(sessionID string) string { t.Helper() var value string - err := pool.QueryRow(t.Context(), `SELECT jsonb_build_object( + err := db.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), @@ -171,8 +171,8 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { if reopen { stop(false) server.Close() - pool.Close() - s, pool = store.NewModelTestStore(t) + db.pool.Close() + s, db = newModelTestStoreDB(t) server, stop = serve() settings["base"] = server.URL } diff --git a/services/core/internal/store/self_hosted_initial_public_test.go b/services/core/internal/store/self_hosted_initial_public_test.go index 896549d7..b8776e53 100644 --- a/services/core/internal/store/self_hosted_initial_public_test.go +++ b/services/core/internal/store/self_hosted_initial_public_test.go @@ -29,7 +29,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, pool := store.NewModelTestStore(t) + s, db := newModelTestStoreDB(t) tenant, foreignTenant := uuid.NewString(), uuid.NewString() token, peer, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ @@ -38,7 +38,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.Store, worker *execution.Worker) *httptest.Server { + serve := func(s *store.Store, db fixtureDB, worker *execution.Worker) *httptest.Server { t.Helper() enabled := []func(*api.Dependencies){acceptUnavailable(t)} if worker != nil { @@ -49,7 +49,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { d.Execution = &api.Execution{ExecutorURL: origin, Admission: unavailableAdmission{}, SessionArchive: strictStandIn{t}, Workspaces: strictStandIn{t}} }) } - handler, err := publicHandler(t, s, auth, "codex", enabled...) + handler, err := publicHandler(t, s, db, auth, "codex", enabled...) if err != nil { t.Fatal(err) } @@ -57,8 +57,8 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { t.Cleanup(server.Close) return server } - worker, stop := publicInitialWorker(t, s) - server := serve(s, worker) + worker, stop := publicInitialWorker(t, s, db) + server := serve(s, db, 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() @@ -137,7 +137,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { } return result } - before := reservations(s, pool) + before := reservations(s, db.pool) for _, reservation := range before { if reservation.State != store.EnvironmentInputPending || reservation.Deadline.Sub(reservation.CreatedAt) != 5*time.Minute { t.Fatal("public initial creation did not retain its database deadline") @@ -145,13 +145,13 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { } stop(false) server.Close() - pool.Close() - reopened, reopenedPool := store.NewModelTestStore(t) - worker, stop = publicInitialWorker(t, reopened) - server = serve(reopened, worker) + db.pool.Close() + reopened, reopenedDB := newModelTestStoreDB(t) + worker, stop = publicInitialWorker(t, reopened, reopenedDB) + server = serve(reopened, reopenedDB, worker) settings["base"], settings["accepted"] = server.URL, accepted run("reopen") - if !reflect.DeepEqual(before, reservations(reopened, reopenedPool)) { + if !reflect.DeepEqual(before, reservations(reopened, reopenedDB.pool)) { t.Fatal("reopened public retry changed reservation identity or deadline") } failureID := created.Cases[0].ID @@ -161,7 +161,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { return } // Advance one known deadline; the running Worker still owns settlement and events. - tag, err := reopenedPool.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 := 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) if err != nil || tag.RowsAffected() != 1 { t.Error("controlled initial deadline update failed", err) w.WriteHeader(http.StatusInternalServerError) @@ -172,7 +172,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { defer control.Close() settings["expiry_control"] = control.URL run("expire") - after := reservations(reopened, reopenedPool) + after := reservations(reopened, reopenedDB.pool) for id, reservation := range after { if id == failureID { if reservation.ID != before[id].ID || reservation.State != store.EnvironmentInputExpired || reservation.SettledAt == nil { @@ -183,7 +183,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { } } var pid uint32 - err := reopenedPool.QueryRow(t.Context(), `SELECT pid FROM pg_locks WHERE locktype='advisory' + err := reopenedDB.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) @@ -191,24 +191,20 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { t.Fatal(err) } var killed bool - if err := reopenedPool.QueryRow(t.Context(), "SELECT pg_terminate_backend($1, 1000)", pid).Scan(&killed); err != nil || !killed { + if err := reopenedDB.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, nil).URL + settings["disabled_base"] = serve(reopened, reopenedDB, nil).URL run("unavailable") - if !reflect.DeepEqual(after, reservations(reopened, reopenedPool)) { + if !reflect.DeepEqual(after, reservations(reopened, reopenedDB.pool)) { t.Fatal("unavailable execution or recorded retry changed initial work") } } -func publicInitialWorker(t *testing.T, s *store.Store) (*execution.Worker, func(bool)) { +func publicInitialWorker(t *testing.T, s *store.Store, db fixtureDB) (*execution.Worker, func(bool)) { t.Helper() - dispatcher := &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()} - worker, err := execution.StartWorker(t.Context(), dispatcher) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()}) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() diff --git a/services/core/internal/store/session_agent_filter_public_test.go b/services/core/internal/store/session_agent_filter_public_test.go index 3e51233f..08fbc257 100644 --- a/services/core/internal/store/session_agent_filter_public_test.go +++ b/services/core/internal/store/session_agent_filter_public_test.go @@ -18,20 +18,20 @@ func TestSessionAgentFilterOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, pool := store.NewTestStore(t) + s, db := newTestStoreDB(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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() - recoveredStore := store.New(pool) - h, err = publicHandler(t, recoveredStore, auth, "codex", storeExecution(t, recoveredStore)) + recoveredStore, recoveredDB := store.New(db.pool), fixtureDB{pool: db.pool} + h, err = publicHandler(t, recoveredStore, recoveredDB, auth, "codex", storeExecution(t, recoveredStore)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_artifacts_public_test.go b/services/core/internal/store/session_artifacts_public_test.go index 7a7cbf9c..704eb404 100644 --- a/services/core/internal/store/session_artifacts_public_test.go +++ b/services/core/internal/store/session_artifacts_public_test.go @@ -68,7 +68,7 @@ func completeArtifactTurn(t *testing.T, s *store.Store, tenant, session, environ } // artifactHTTPServer serves Artifact routes for an owner and a foreign tenant. -func artifactHTTPServer(t *testing.T, s *store.Store) (server *httptest.Server, owner, ownerTenant, foreign, foreignTenant string) { +func artifactHTTPServer(t *testing.T, s *store.Store, db fixtureDB) (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.Store) (server *httptest.Server, {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, auth, "codex") + h, err := publicHandler(t, s, db, auth, "codex") if err != nil { t.Fatal(err) } @@ -89,8 +89,8 @@ func artifactHTTPServer(t *testing.T, s *store.Store) (server *httptest.Server, // environment_id filter matches nothing like another Environment's ID (HE-56), // without weakening tenant or Session scoping. func TestSessionArtifactListEnvelopeAndEnvironmentFilterPostgres(t *testing.T) { - s, _ := store.NewTestStore(t) - server, owner, ownerTenant, foreign, foreignTenant := artifactHTTPServer(t, s) + s, db := newTestStoreDB(t) + server, owner, ownerTenant, foreign, foreignTenant := artifactHTTPServer(t, s, db) client := pathIDClient{t: t, server: server} session, environment := hostedArtifactSession(t, s, ownerTenant, "artifact-list") @@ -209,8 +209,8 @@ func TestSessionArtifactsOfficialClientPostgres(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, _ := store.NewTestStore(t) - server, owner, ownerTenant, foreign, _ := artifactHTTPServer(t, s) + s, db := newTestStoreDB(t) + server, owner, ownerTenant, foreign, _ := artifactHTTPServer(t, s, db) 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, ownerTenant, session, environment, "artifact-sdk-1", outputs) diff --git a/services/core/internal/store/session_deletion_execution_test.go b/services/core/internal/store/session_deletion_execution_test.go index 13c1dcfe..8f1e138a 100644 --- a/services/core/internal/store/session_deletion_execution_test.go +++ b/services/core/internal/store/session_deletion_execution_test.go @@ -8,7 +8,6 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) @@ -16,10 +15,7 @@ func TestDeletedSessionWaitingTurnSettlesWithoutStoppingWorker(t *testing.T) { h := newFunctionHarness(t) ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -81,10 +77,7 @@ func TestDeletedSessionRestartStillReconcilesHiddenClaim(t *testing.T) { if err := h.s.CommitLegacyDeletion(ctx, h.tenant, h.session.ID); err != nil { t.Fatal(err) } - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) stopped, cancel := context.WithCancel(ctx) cancel() if err := worker.Run(stopped); !errors.Is(err, context.Canceled) { @@ -105,10 +98,7 @@ func TestWaitingSessionCancelsThenDeletesThroughWorker(t *testing.T) { h := newFunctionHarness(t) ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { diff --git a/services/core/internal/store/session_deletion_lifecycle_public_test.go b/services/core/internal/store/session_deletion_lifecycle_public_test.go index 75e15b00..9e2d1c1f 100644 --- a/services/core/internal/store/session_deletion_lifecycle_public_test.go +++ b/services/core/internal/store/session_deletion_lifecycle_public_test.go @@ -1,13 +1,11 @@ package store_test import ( - "context" "encoding/json" "net/http" "net/http/httptest" "reflect" "testing" - "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" @@ -22,29 +20,21 @@ 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, pool := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(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: "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, auth, "codex", storeExecution(t, s), executorURL("https://executor.example")) + h, err := publicHandler(t, s, db, 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} - writer, err := store.NewExecution(ctx, s) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { - closing, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - _ = writer.CloseExecution(closing) - }) + writer := executionOwner(t, db, s).Store create := func(environment string, initial bool) store.Session { t.Helper() @@ -140,7 +130,7 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { "later_input_cancelled": "", } expired := create(selfHosted, true) - if _, err := pool.Exec(ctx, "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1", expired.ID); err != nil { + 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 { t.Fatal(err) } if count, err := writer.ExpireEnvironmentInputs(ctx); err != nil || count != 1 { @@ -187,13 +177,13 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { if readStatus != http.StatusOK || decode(before)["status"] != projected[name] { t.Fatal(readStatus, before) } - digest := databaseDigest(t, pool) + digest := databaseDigest(t, db.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, pool); !reflect.DeepEqual(after, digest) { + if after := databaseDigest(t, db.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 { @@ -208,7 +198,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, pool) + digest := databaseDigest(t, db.pool) filter := store.WriteOperationFilter{ResourceType: "session", ResourceID: id, Limit: 100} beforeAudit, err := s.ListWriteOperations(ctx, tenant, filter) if err != nil { @@ -230,7 +220,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, pool) + after := databaseDigest(t, db.pool) delete(after, "write_audit_operations") delete(digest, "write_audit_operations") if !reflect.DeepEqual(after, digest) { diff --git a/services/core/internal/store/session_deletion_public_test.go b/services/core/internal/store/session_deletion_public_test.go index fcae8aea..ca34c912 100644 --- a/services/core/internal/store/session_deletion_public_test.go +++ b/services/core/internal/store/session_deletion_public_test.go @@ -18,19 +18,19 @@ func TestSessionDeletionOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, pool := store.NewTestStore(t) + s, db := newTestStoreDB(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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() - h, err = publicHandler(t, store.New(pool), auth, "codex", storeExecution(t, store.New(pool))) + h, err = publicHandler(t, store.New(db.pool), db, auth, "codex", storeExecution(t, store.New(db.pool))) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_initial_public_test.go b/services/core/internal/store/session_initial_public_test.go index 85f4a973..da77de17 100644 --- a/services/core/internal/store/session_initial_public_test.go +++ b/services/core/internal/store/session_initial_public_test.go @@ -9,7 +9,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -18,14 +17,11 @@ func TestInitialSessionInputOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, _ := store.NewTestStore(t) + s, db := newTestStoreDB(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, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s}) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s}) t.Cleanup(func() { stopped, cancel := context.WithCancel(context.Background()) cancel() @@ -33,13 +29,13 @@ func TestInitialSessionInputOfficialClient(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, auth, "codex", workerExecution(worker)) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(worker)) if err != nil { t.Fatal(err) } server := httptest.NewServer(handler) defer server.Close() - unsupported, err := publicHandler(t, s, auth, "fake_alpha", workerExecution(worker)) + unsupported, err := publicHandler(t, s, db, auth, "fake_alpha", workerExecution(worker)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_model_execution_http_test.go b/services/core/internal/store/session_model_execution_http_test.go index c5342f3e..86ee0f46 100644 --- a/services/core/internal/store/session_model_execution_http_test.go +++ b/services/core/internal/store/session_model_execution_http_test.go @@ -16,10 +16,10 @@ import ( func TestModelExecutionHTTPWriteOnlyAndStrictAdmission(t *testing.T) { _, pool := store.NewTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{6}, 32)) - st := store.NewWithCredentialCipher(pool, cipher) + st, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, st), managedSandboxes(t, st)) + handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, st)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_reference_retry_public_test.go b/services/core/internal/store/session_reference_retry_public_test.go index 0e245ced..c2068948 100644 --- a/services/core/internal/store/session_reference_retry_public_test.go +++ b/services/core/internal/store/session_reference_retry_public_test.go @@ -21,13 +21,10 @@ func TestSavedReferenceRetryOfficialClient(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, pool := store.NewTestStore(t) + s, db := newTestStoreDB(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, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s}) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s}) t.Cleanup(func() { ctx, cancel := context.WithCancel(context.Background()) cancel() @@ -35,13 +32,13 @@ func TestSavedReferenceRetryOfficialClient(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, auth, "codex", workerExecution(worker)) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(worker)) if err != nil { t.Fatal(err) } server := httptest.NewServer(handler) defer server.Close() - recovered, err := publicHandler(t, store.New(pool), auth, "codex") + recovered, err := publicHandler(t, store.New(db.pool), db, auth, "codex") if err != nil { t.Fatal(err) } @@ -60,9 +57,9 @@ func TestSavedReferenceRetryOfficialClient(t *testing.T) { } var err error if input.Delete { - _, err = pool.Exec(r.Context(), `DELETE FROM agents WHERE tenant_id=$1 AND id=$2`, tenant, input.ID) + _, err = db.pool.Exec(r.Context(), `DELETE FROM agents WHERE tenant_id=$1 AND id=$2`, tenant, input.ID) } else { - _, err = pool.Exec(r.Context(), `UPDATE agents SET configuration=configuration || $3::jsonb WHERE tenant_id=$1 AND id=$2`, tenant, input.ID, input.Patch) + _, 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) } if err != nil { t.Error(err) diff --git a/services/core/internal/store/sessions.go b/services/core/internal/store/sessions.go index a4a2e86b..47f32f83 100644 --- a/services/core/internal/store/sessions.go +++ b/services/core/internal/store/sessions.go @@ -84,13 +84,13 @@ type SessionPage struct { type Store struct { queries *sqlc.Queries - // pool supplies the execution lease's dedicated connection; pooled runs - // every other transaction. + // pool is the database the Store was built on; pooled runs every + // transaction that is not on the execution lease. pool *pgxpool.Pool pooled *pgunit.Pool // writer runs Session and execution-only transactions, and lease grants // execution authority. New sets writer to pooled and leaves lease nil; - // NewExecution sets both to the same execution lease. Neither changes later. + // NewExecution sets both to the lease it borrows. Neither changes later. writer transactor lease *pgunit.Lease credentialCipher *credentialcrypto.Cipher diff --git a/services/core/internal/store/skill_selectors_public_test.go b/services/core/internal/store/skill_selectors_public_test.go index 2a8c1ad3..f6758fbd 100644 --- a/services/core/internal/store/skill_selectors_public_test.go +++ b/services/core/internal/store/skill_selectors_public_test.go @@ -26,20 +26,20 @@ func TestSkillSelectorsOfficialClientPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex") + h, err := publicHandler(t, s, db, auth, "codex") if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() - recoveredStore := store.NewWithCredentialCipher(pool, cipher) - h, err = publicHandler(t, recoveredStore, auth, "codex") + recoveredStore, recoveredDB := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + h, err = publicHandler(t, recoveredStore, recoveredDB, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/skill_version_deletion_public_test.go b/services/core/internal/store/skill_version_deletion_public_test.go index 3d05c650..d63940cd 100644 --- a/services/core/internal/store/skill_version_deletion_public_test.go +++ b/services/core/internal/store/skill_version_deletion_public_test.go @@ -25,13 +25,13 @@ func TestSkillVersionDeletionHTTPPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex") + h, err := publicHandler(t, s, db, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/skills_public_test.go b/services/core/internal/store/skills_public_test.go index 3b42c7a0..98e6ee95 100644 --- a/services/core/internal/store/skills_public_test.go +++ b/services/core/internal/store/skills_public_test.go @@ -26,20 +26,20 @@ func TestSkillsOfficialClientPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex") + h, err := publicHandler(t, s, db, auth, "codex") if err != nil { t.Fatal(err) } server := httptest.NewServer(h) defer server.Close() - recoveredStore := store.NewWithCredentialCipher(pool, cipher) - h, err = publicHandler(t, recoveredStore, auth, "codex") + recoveredStore, recoveredDB := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + h, err = publicHandler(t, recoveredStore, recoveredDB, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/source_files_errors_public_test.go b/services/core/internal/store/source_files_errors_public_test.go index 51ae3005..9fd35f39 100644 --- a/services/core/internal/store/source_files_errors_public_test.go +++ b/services/core/internal/store/source_files_errors_public_test.go @@ -9,7 +9,6 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -18,13 +17,13 @@ func TestSourceFileErrorsOfficialClientPostgres(t *testing.T) { if python == "" { t.Skip("pinned official Python SDK required") } - s, _ := store.NewTestStore(t) + s, db := newTestStoreDB(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, auth, "codex") + h, err := publicHandler(t, s, db, auth, "codex") if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/stream_authority_http_test.go b/services/core/internal/store/stream_authority_http_test.go index 98161177..782e617d 100644 --- a/services/core/internal/store/stream_authority_http_test.go +++ b/services/core/internal/store/stream_authority_http_test.go @@ -21,7 +21,7 @@ func TestLiveStreamClosesAfterKeyRevocationOrProjectArchive(t *testing.T) { name = "project-archive" } t.Run(name, func(t *testing.T) { - s, _ := store.NewTestStore(t) + s, db := newTestStoreDB(t) projectID := uuid.NewString() ctx := adminaudit.WithSource(t.Context(), adminaudit.Source{CredentialID: "12345678", ActorLabel: "test", RequestID: uuid.NewString(), TraceID: uuid.NewString(), ProjectID: projectID}) project, err := s.CreateProject(ctx, projectID, "Stream authority") @@ -40,7 +40,7 @@ func TestLiveStreamClosesAfterKeyRevocationOrProjectArchive(t *testing.T) { if err != nil { t.Fatal(err) } - h, err := publicHandler(t, s, nil, "codex", storeKeys(s), storeExecution(t, s)) + h, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), storeExecution(t, s)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/structured_output_native_test.go b/services/core/internal/store/structured_output_native_test.go index c9a313fa..9d2149b9 100644 --- a/services/core/internal/store/structured_output_native_test.go +++ b/services/core/internal/store/structured_output_native_test.go @@ -28,10 +28,7 @@ func TestNativeStructuredOutputPublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -47,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, auth, "claude_sdk", workerExecution(worker), nativeDeploymentDefaults(h.s, model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "claude_sdk", workerExecution(worker), nativeDeploymentDefaults(h.s, model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/subagent_dispatch_test.go b/services/core/internal/store/subagent_dispatch_test.go index 0b955f7e..a2cc7555 100644 --- a/services/core/internal/store/subagent_dispatch_test.go +++ b/services/core/internal/store/subagent_dispatch_test.go @@ -1,7 +1,6 @@ package store_test import ( - "context" "encoding/json" "errors" "testing" @@ -27,12 +26,7 @@ func TestSubagentIdentityUsesLeasedDispatchJournal(t *testing.T) { if err = h.s.BindSessionDevice(ctx, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } - writer, err := store.NewExecution(ctx, h.s) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = writer.CloseExecution(context.Background()) }) - h.d.Store = writer + h.d.Store = executionOwner(t, h.db, h.s).Store input := h.message("first", "root message") running := h.run(ctx, input.TurnID) var request proto.PromptRequestPayload diff --git a/services/core/internal/store/subagent_identities_test.go b/services/core/internal/store/subagent_identities_test.go index 06a046f3..6eac9f0d 100644 --- a/services/core/internal/store/subagent_identities_test.go +++ b/services/core/internal/store/subagent_identities_test.go @@ -109,7 +109,7 @@ func TestSubagentIdentityIsAtomicScopedAndImmutable(t *testing.T) { if _, err = w.CompleteExecution(ctx, tenant, session.ID, input.TurnID, TurnCompleted, json.RawMessage(`{}`), "root", input.Sequence); err != nil { t.Fatal(err) } - if err = w.CloseExecution(ctx); err != nil { + if err = w.lease.Close(ctx); err != nil { t.Fatal(err) } reopened, _ := testStore(t) @@ -157,7 +157,7 @@ func TestSubagentIdentityRejectsLostLease(t *testing.T) { if err := old.AppendTurnEvents(t.Context(), tenant, session.ID, input.TurnID, 1, []ExecutionEvent{subagentIdentityEvent("child", "root", 100)}); err == nil { t.Fatal("lost owner committed identity") } - if err := successor.CheckExecutionOwnership(context.Background()); err != nil { + if err := successor.lease.CheckOwnership(context.Background()); err != nil { t.Fatal(err) } events, err := s.ListTurnEvents(t.Context(), tenant, session.ID, input.TurnID, 0, 100) diff --git a/services/core/internal/store/subagent_visibility_public_test.go b/services/core/internal/store/subagent_visibility_public_test.go index ae2d33ae..d77da801 100644 --- a/services/core/internal/store/subagent_visibility_public_test.go +++ b/services/core/internal/store/subagent_visibility_public_test.go @@ -69,24 +69,20 @@ func subagentFixture(kind string, value any) store.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, _ := store.NewTestStore(t) + s, db := newTestStoreDB(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, auth, "codex", storeExecution(t, s)) + handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } server := httptest.NewServer(handler) defer server.Close() client := pathIDClient{t: t, server: server} - writer, err := store.NewExecution(t.Context(), s) - if err != nil { - t.Fatal(err) - } - defer func() { _ = writer.CloseExecution(t.Context()) }() + writer := executionOwner(t, db, s).Store ctx := t.Context() created := openStream(t, server, token, http.MethodPost, "/v1/agents/sessions", diff --git a/services/core/internal/store/template_composition_public_test.go b/services/core/internal/store/template_composition_public_test.go index a27d8869..4aee834b 100644 --- a/services/core/internal/store/template_composition_public_test.go +++ b/services/core/internal/store/template_composition_public_test.go @@ -30,8 +30,8 @@ func TestTemplateCompositionOfficialClientPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) - reopenedStore := store.NewWithCredentialCipher(pool, cipher) + s, reopenedStore := store.NewWithCredentialCipher(pool, cipher), store.NewWithCredentialCipher(pool, cipher) + db := fixtureDB{pool: pool, cipher: cipher} // built both Stores 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}, @@ -40,7 +40,7 @@ func TestTemplateCompositionOfficialClientPostgres(t *testing.T) { serve := func(current *store.Store) *httptest.Server { t.Helper() // Hosted admission and freezing use the real Store; no Runtime or model runs. - h, err := publicHandler(t, current, auth, "codex", storeExecution(t, current), managedSandboxes(t, current), fixtureDeploymentProvider(current)) + h, err := publicHandler(t, current, db, auth, "codex", storeExecution(t, current), managedSandboxes(t, current), fixtureDeploymentProvider(current)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/template_null_selection_public_test.go b/services/core/internal/store/template_null_selection_public_test.go index 48d31b75..7156f74b 100644 --- a/services/core/internal/store/template_null_selection_public_test.go +++ b/services/core/internal/store/template_null_selection_public_test.go @@ -32,8 +32,8 @@ func TestTemplateNullSelectionOfficialClientPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) - reopenedStore := store.NewWithCredentialCipher(pool, cipher) + s, reopenedStore := store.NewWithCredentialCipher(pool, cipher), store.NewWithCredentialCipher(pool, cipher) + db := fixtureDB{pool: pool, cipher: cipher} // built both Stores 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}, @@ -41,7 +41,7 @@ func TestTemplateNullSelectionOfficialClientPostgres(t *testing.T) { }) serve := func(current *store.Store) *httptest.Server { t.Helper() - h, err := publicHandler(t, current, auth, "codex", storeExecution(t, current), managedSandboxes(t, current), fixtureDeploymentProvider(current)) + h, err := publicHandler(t, current, db, auth, "codex", storeExecution(t, current), managedSandboxes(t, current), fixtureDeploymentProvider(current)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/tool_policy_native_test.go b/services/core/internal/store/tool_policy_native_test.go index 141612ac..35e02579 100644 --- a/services/core/internal/store/tool_policy_native_test.go +++ b/services/core/internal/store/tool_policy_native_test.go @@ -34,10 +34,7 @@ func TestNativeToolPolicyPublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 15*time.Minute) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -53,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, auth, kind, workerExecution(worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(h.s, model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(h.s, model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/tool_search_native_test.go b/services/core/internal/store/tool_search_native_test.go index 33ec2eee..c36fe169 100644 --- a/services/core/internal/store/tool_search_native_test.go +++ b/services/core/internal/store/tool_search_native_test.go @@ -28,10 +28,7 @@ func TestNativeToolSearchPublicExecution(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, h.db, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -47,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, auth, "claude_sdk", workerExecution(worker), nativeDeploymentDefaults(h.s, model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "claude_sdk", workerExecution(worker), nativeDeploymentDefaults(h.s, model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/unified_model_configuration_http_test.go b/services/core/internal/store/unified_model_configuration_http_test.go index f1e6787f..cfed60a7 100644 --- a/services/core/internal/store/unified_model_configuration_http_test.go +++ b/services/core/internal/store/unified_model_configuration_http_test.go @@ -10,19 +10,18 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/api" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) func TestUnifiedModelConfigurationHTTP(t *testing.T) { - st, _ := store.NewManagedTestStore(t) + st, db := newManagedTestStoreDB(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, auth, "codex", storeExecution(t, st), managedSandboxes(t, st), withCoreKeys(admin), withHarnesses([]string{"codex", "claude_sdk"})) + handler, err := publicHandler(t, st, db, 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/internal/store/unstorable_text_public_test.go b/services/core/internal/store/unstorable_text_public_test.go index f389bc08..fb611768 100644 --- a/services/core/internal/store/unstorable_text_public_test.go +++ b/services/core/internal/store/unstorable_text_public_test.go @@ -25,10 +25,10 @@ func TestUnstorableTextRejectsWithoutWritesPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: 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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/whitespace_input_public_test.go b/services/core/internal/store/whitespace_input_public_test.go index 1fb64838..d03252e9 100644 --- a/services/core/internal/store/whitespace_input_public_test.go +++ b/services/core/internal/store/whitespace_input_public_test.go @@ -11,7 +11,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -20,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, pool := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(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, auth, "codex", storeExecution(t, s)) + h, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s)) if err != nil { t.Fatal(err) } @@ -73,7 +72,7 @@ func TestWhitespaceInputStoredVerbatimPostgres(t *testing.T) { } // W4: unchanged rejection without writes. - before := databaseDigest(t, pool) + before := databaseDigest(t, db.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 { @@ -85,7 +84,7 @@ func TestWhitespaceInputStoredVerbatimPostgres(t *testing.T) { t.Errorf("events %s: %d %s", input, status, body) } } - if after := databaseDigest(t, pool); !mapsEqual(before, after) { + if after := databaseDigest(t, db.pool); !mapsEqual(before, after) { t.Error("rejected empty input changed persisted state") } } @@ -94,14 +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, pool := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(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, err := execution.StartWorker(t.Context(), &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()}) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()}) t.Cleanup(func() { stopped, cancel := context.WithCancel(context.Background()) cancel() @@ -110,7 +106,7 @@ func TestWhitespaceOnlyTextHarnessAdmissionPostgres(t *testing.T) { } }) serve := func(engine string) pathIDClient { - handler, err := publicHandler(t, s, auth, engine, workerExecution(worker), executorURL("https://offline-executor.example")) + handler, err := publicHandler(t, s, db, auth, engine, workerExecution(worker), executorURL("https://offline-executor.example")) if err != nil { t.Fatal(err) } @@ -143,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, pool) + before := databaseDigest(t, db.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"}]}]}`, @@ -161,7 +157,7 @@ func TestWhitespaceOnlyTextHarnessAdmissionPostgres(t *testing.T) { t.Errorf("%s events %s: %d %s", engine, body, status, response) } } - if after := databaseDigest(t, pool); !mapsEqual(before, after) { + if after := databaseDigest(t, db.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/internal/store/worker_capacity_test.go b/services/core/internal/store/worker_capacity_test.go index a81840ff..fe02256e 100644 --- a/services/core/internal/store/worker_capacity_test.go +++ b/services/core/internal/store/worker_capacity_test.go @@ -42,7 +42,7 @@ func TestWorkerDefersPreparationCapacityUntilCleanupReleasesSlot(t *testing.T) { sessions[h.session.ID] = receipt.TurnID } frames := capacityWorkerFrames(t, h) - _, stop := startEnvironmentExpiryWorker(t, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, 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.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) defer stop() prepare := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) if test.operation == proto.TypeExecutionStart { diff --git a/services/core/internal/store/worker_input_race_test.go b/services/core/internal/store/worker_input_race_test.go index c4a3953a..7dba1f1e 100644 --- a/services/core/internal/store/worker_input_race_test.go +++ b/services/core/internal/store/worker_input_race_test.go @@ -9,7 +9,6 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" @@ -61,10 +60,7 @@ func TestWorkerInputReadSkipsConcurrentlyCancelledCandidate(t *testing.T) { h.d.Store = store.New(instrumented) ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, fixtureDB{pool: instrumented}, h.d) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { diff --git a/services/core/internal/store/worker_lease_loss_test.go b/services/core/internal/store/worker_lease_loss_test.go index 05f75ff5..ad420eea 100644 --- a/services/core/internal/store/worker_lease_loss_test.go +++ b/services/core/internal/store/worker_lease_loss_test.go @@ -8,7 +8,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) @@ -18,10 +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, err := execution.StartWorker(t.Context(), h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, t.Context(), h.db, h.d) ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() @@ -75,10 +71,7 @@ func TestWorkerLeaseLossLeavesUncertainWorkForSuccessor(t *testing.T) { if err != nil || active.Status != store.TurnInProgress { t.Fatal("lost owner persisted fallback completion", active, err) } - successor, err := execution.StartWorker(t.Context(), h.d) - if err != nil { - t.Fatal(err) - } + successor := startWorker(t, t.Context(), h.db, h.d) stopped, stop := context.WithCancel(t.Context()) stop() if err = successor.Run(stopped); err != context.Canceled { diff --git a/services/core/internal/store/worker_preparation_failure_test.go b/services/core/internal/store/worker_preparation_failure_test.go index 6ce1549b..fb1b2239 100644 --- a/services/core/internal/store/worker_preparation_failure_test.go +++ b/services/core/internal/store/worker_preparation_failure_test.go @@ -19,7 +19,7 @@ func TestWorkerSettlesConfirmedPreparationFailureAndAcceptsNewInput(t *testing.T if err != nil { t.Fatal(err) } - _, stop := startEnvironmentExpiryWorker(t, h.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, 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.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, 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.d) + _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) defer stop() old := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) if _, err := h.s.CancelEnvironmentInput(t.Context(), h.tenant, h.session.ID, first.ID); err != nil { diff --git a/services/core/internal/store/worker_wakeup_test.go b/services/core/internal/store/worker_wakeup_test.go index 64e5478f..a8be6fcc 100644 --- a/services/core/internal/store/worker_wakeup_test.go +++ b/services/core/internal/store/worker_wakeup_test.go @@ -9,7 +9,6 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" @@ -48,10 +47,7 @@ func TestWorkerSchedulerCommittedAdmissionWakesBeforeMaintenance(t *testing.T) { h.d.Store = store.NewWithCredentialCipher(instrumented, store.FixtureCipher()) ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker, err := execution.StartWorker(ctx, h.d) - if err != nil { - t.Fatal(err) - } + worker := startWorker(t, ctx, fixtureDB{pool: instrumented, cipher: store.FixtureCipher()}, h.d) done := make(chan error, 1) started := false defer func() { @@ -112,7 +108,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.d) + worker, stop := startEnvironmentExpiryWorker(t, h.db, h.d) defer stop() awaitDaemonRemoteCondition(t, t.Context(), 5*time.Second, "initial empty scheduler scan", func() bool { return worker.MetricsSnapshot().Scheduler.LastRunAt != nil