diff --git a/contracts/agents-api/core.openapi.yaml b/contracts/agents-api/core.openapi.yaml index 6699d80a6..f15a93129 100644 --- a/contracts/agents-api/core.openapi.yaml +++ b/contracts/agents-api/core.openapi.yaml @@ -546,7 +546,7 @@ definitions: properties: data: items: - $ref: '#/definitions/store.RuntimeNodeAllocation' + $ref: '#/definitions/deployment.NodeAllocation' type: array type: object api.SandboxDeploymentChangeInput: @@ -998,6 +998,35 @@ definitions: snapshots: type: integer type: object + deployment.NodeAllocation: + properties: + compute_phase: + type: string + compute_phase_changed_at: + description: The time the allocation entered its current compute_phase, or null when unknown; an allocation that existed before Core recorded it reports null until its next phase change. For a suspended microsandbox allocation, this time plus the deployment's snapshot retention tells roughly when Core reclaims it. + type: string + x-nullable: true + created_at: + type: string + deployment_generation: + type: integer + diagnostic: + type: string + environment_id: + type: string + id: + type: string + initialization: + type: string + node_id: + type: string + session_id: + type: string + state: + type: string + tenant_id: + type: string + type: object deployment.NodeDetail: properties: active: @@ -1408,35 +1437,6 @@ definitions: vaults: type: integer type: object - store.RuntimeNodeAllocation: - properties: - compute_phase: - type: string - compute_phase_changed_at: - description: The time the allocation entered its current compute_phase, or null when unknown; an allocation that existed before Core recorded it reports null until its next phase change. For a suspended microsandbox allocation, this time plus the deployment's snapshot retention tells roughly when Core reclaims it. - type: string - x-nullable: true - created_at: - type: string - deployment_generation: - type: integer - diagnostic: - type: string - environment_id: - type: string - id: - type: string - initialization: - type: string - node_id: - type: string - session_id: - type: string - state: - type: string - tenant_id: - type: string - type: object v1.Agent: properties: id: diff --git a/packages/agents-client/src/sandbox-client.test.ts b/packages/agents-client/src/sandbox-client.test.ts index ec61820b4..7567263b4 100644 --- a/packages/agents-client/src/sandbox-client.test.ts +++ b/packages/agents-client/src/sandbox-client.test.ts @@ -6,7 +6,7 @@ import nodeDiagnosticFixture from "../../../services/core/internal/sandbox/testd function response(value: unknown, status = 200) { return new Response(JSON.stringify(value), { status }); } -// Shapes as Core serializes them (deployment.Node, deployment.NodeDetail, store.RuntimeNodeAllocation, deployment.View). +// Shapes as Core serializes them (deployment.Node, deployment.NodeDetail, deployment.NodeAllocation, deployment.View). const created = "2026-09-25T08:00:00.123456789Z"; /** A ready node: Core omits `diagnostic`. */ const node = { diff --git a/packages/agents-client/src/sandbox-client.ts b/packages/agents-client/src/sandbox-client.ts index 47269b1ad..ef18d0c02 100644 --- a/packages/agents-client/src/sandbox-client.ts +++ b/packages/agents-client/src/sandbox-client.ts @@ -180,7 +180,7 @@ function invalidSandboxResponse(): never { /** * The object has every required member, optional ones only where listed, and nothing else. The projections below - * follow Core's deployment.View, deployment.Node, deployment.NodeDetail and store.RuntimeNodeAllocation serialization. + * follow Core's deployment.View, deployment.Node, deployment.NodeDetail and deployment.NodeAllocation serialization. */ function members(value: unknown, required: readonly string[], optional: readonly string[] = []): Record { if (!isRecord(value) || !required.every((field) => hasOwn(value, field)) || !onlyFields(value, new Set([...required, ...optional]))) return invalidSandboxResponse(); diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 3d6aac68a..ee1ceef71 100644 --- a/services/core/IMPLEMENTATION.md +++ b/services/core/IMPLEMENTATION.md @@ -18,7 +18,7 @@ Shared vocabulary has one owner each, and domains use it rather than copy it. `i `internal/sessions` owns the Session vocabulary: Sessions, Turns, inputs, Environments and their provisioning failures, function calls, Item and Artifact reads, executor credentials, and the errors Session operations return, which `api` maps in `writeSessionsError`. It also owns the Session change vocabulary and its decisions: the public changes that report Turn and Session transitions, what a Turn that ends settles, measured Turn usage and the Session activity a change reports. Sequences that several Session operations share are `sessions` procedures, each over the transaction interface it declares: `CancelTurn` and `CancelWork` cancel work, `TrackInputActivity` reports the input activity a write changes, `FailEnvironment` and `TerminateEnvironment` settle an Environment that failed or whose managed compute ended, `CreateEnvironmentDevice` creates and binds a hosted Environment's dedicated Runtime device, `CheckComputeAdmission`, `CheckFileWriteGate` and `CheckInputStart` gate new work, `AdmitFunctionResult` saves a submitted function result for its Turn's call, and `LoadRequiredActions` lists the calls a Turn waits on. `AppendTurnEvents` records a batch of a Turn's execution observations, which `NewJournalBatch` validates, and `AppendTurnEvent` records its terminal outcome; both project the new entries. `ProjectSource` projects one journal entry or admitted input into Items, the root Turn's usage and Subagents, and `ProjectInput` projects an admitted input. `ValidTransition` decides the Turn status transitions, and the Subagent projection decides each Subagent's identity, lifecycle, child Turn and child Item changes. `sessions.ExecutionOperations` runs the Session writes only the execution owner makes; each function-call operation loads the Turn and its call in one leased Session transaction, decides in `sessions`, then applies in `sessionpg`, and `AppendTurnEvents` records and projects a Turn's journal batch in one leased Session transaction. `internal/items` owns public Items: it projects observations, merges them into stored Items and builds the ordered events that report each Item change. Neither imports persistence. `internal/persistence/postgres/sessionpg` loads the facts those decisions read and applies them inside the caller's Session transaction, under the Session lock: it allocates event sequence positions, event IDs, Item positions, output indexes and Subagent IDs, writes the change journal, each Turn's execution journal and its counters, Items, Turn usage, Subagent bindings, child Turns and child Items, and Artifact settlement, and prunes the change journal. It decides nothing. `WithSession` runs a Session operation's transaction on the pool or the lease: it locks the Session, applies the operation, then prunes the change journal. `BindSession` binds one Session to a transaction's queries as a `SessionTx`, which implements the procedures' interfaces. -Other adapters call only these `sessionpg` participants, on their own transaction's queries: `LockSession`, which orders a write against the Session's Turns; `BindSession`, which runs a Session procedure, such as `CreateEnvironmentDevice`, in that transaction; `LoadEnvironment`, which reads the tenant's Environment of a live Session; and `CreateSessionEnvironment`, which stores the Environment of a Session being created. The other `sessionpg` exports serve `store` until the Session operations leave it. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; deployment's placement fact-loading and reservation code sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. +Other adapters call only these `sessionpg` participants, on their own transaction's queries: `LockSession`, which orders a write against the Session's Turns; `BindSession`, which runs a Session procedure, such as `CreateEnvironmentDevice`, in that transaction; `LoadEnvironment`, which reads the tenant's Environment of a live Session; and `CreateSessionEnvironment`, which stores the Environment of a Session being created. The other `sessionpg` exports serve `store` until the Session operations leave it. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; `placementpg`, which loads the placement facts and reserves placements, sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. `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`. @@ -33,7 +33,7 @@ Domain owners, each with its PostgreSQL adapter under `internal/persistence/post - `environmenttemplates` (`templatepg`): Environment Templates, their validation and default network, their sealed setup, initial files, Skills and Plugins, and the resolved Template that Session creation composes into its Environment. - `modelconfiguration` (`modelconfigurationpg`): each Harness's deployment default model configuration and its last-use observations. - `skills` (`skillpg`): Skills and their immutable versions: archive checks, the default and latest pointers, version selection and deletion, and each version's sealed archive. Session creation freezes selected versions inside its `store` transaction with the `skills` rules. -- `deployment` (`deploymentpg`): the sandbox deployment and its nodes: provider configuration and the sealed credential, the specification and retained generations, setup, update and switch, node enrollment, identity and authentication, generation configuration, capacity, presence, status, host history, reset, and the counts of nodes and sandboxes bound to the public address. It interprets Sandbox Provider declarations through the `providers.Registry` it is given, whose lookups return typed errors. `cmd/server` builds that registry and calls it only to build a direct Provider and to discover a Provider's configuration; it takes each setup's mode and declared operations from the `Setup` that `deployment` returns. The only other reader is `store`, whose own `providers.Builtin()` validates the stored specification at Session admission and reads the public-origin requirement at placement. Deployment changes run through `deployment.ExecutionOperations` on `deploymentpg.NewExecution(lease, …)`, which the Worker receives as `execution.Owner.Deployment`; the Worker reads the deployment and prepares a selection's setup through the pooled `deployment.Service` in `execution.Dispatcher.Deployment` and lists the Sessions a reset still has to archive through the pooled `deployment.Reader` in `execution.Dispatcher.DeploymentReader`, and node management and reads use the pooled `deploymentpg.Store`, which `cmd/server` also reads the owner epoch from and runs the host-history sampler on. Provider calls run outside transactions, and the final transaction rechecks the expected generation. Allocations, placement, reset admission and the reset's Session archive stay in `store`. +- `deployment` and its subpackage `deployment/placement` (`deploymentpg`, `placementpg`): the sandbox deployment and its nodes: provider configuration and the sealed credential, the specification and retained generations, setup, update and switch, node enrollment, identity and authentication, generation configuration, capacity, presence, status, host history, reset, and the counts of nodes and sandboxes bound to the public address; and hosted runtime allocations: reservation with the dedicated device, compute ownership and settlement, observation diagnostics, activity, compute phases, wake receipts and cleanup, with the reads that schedule, discover and authorize them. It interprets Sandbox Provider declarations through the `providers.Registry` it is given, whose lookups return typed errors. `cmd/server` builds that registry and calls it only to build a direct Provider and to discover a Provider's configuration; it takes each setup's mode and declared operations from the `Setup` that `deployment` returns. `deployment/placement` owns hosted admission and placement: `cmd/server` builds one `placement.Rules` from the registry and the public URL, both fixed while Core runs, and gives it to `deployment.Service` and to `store`; its pure decisions admit a hosted Session, choose its node, and admit an allocation's reserved node and a restore on it, and its errors keep one status, code and message through `writeStoreError` and `writeDeploymentError`. `placementpg` loads the facts those decisions read and applies them on the caller's transaction-bound queries; it has no Store or transaction runner. The only other reader is `store`: Session creation decides admission and placement with those rules and `placementpg` inside its own transaction. Deployment changes and allocation writes run through `deployment.ExecutionOperations` on `deploymentpg.NewExecution(lease, …)`, which the Worker receives as `execution.Owner.Deployment`. Each allocation write locks the owning Session, decides in `deployment`, applies in `deploymentpg` and prunes the Session's change journal in the same leased transaction; cleanup settles the Session through the `sessions` procedures on a `sessionpg.SessionTx` bound to that transaction. The Worker reads the deployment, prepares a selection's setup and records live-compute activity through the pooled `deployment.Service` in `execution.Dispatcher.Deployment`, and lists the Sessions a reset still has to archive, schedules node lifecycles and reads allocations through the pooled `deployment.Reader` in `execution.Dispatcher.DeploymentReader`; a pooled read carries no lease, and the leased write that follows it rechecks the allocation's owner. Node management and reads use the pooled `deploymentpg.Store`, which `cmd/server` also reads the owner epoch from and runs the host-history sampler on. Provider calls run outside transactions, and the final transaction rechecks the expected generation. The reset's Session archive stays in `store`. - `sessions` (`sessionpg`): Session use cases and reads as they leave `store`. The pooled `sessions.Service` runs the use cases on `sessionpg.Store`, which implements `sessions.Storage`, and plain reads use `sessions.Reader`, which `sessionpg.Store` also implements, directly. So far these cover root Items and Subagents (their reads), Session Artifacts (their reads, deletion and the staging of a Turn's export), Environments (their reads, the initialization list and the frozen setup and initial files, which `sessionpg.Store` opens with the credential key), devices (their reads, creation, revocation, heartbeats and Runtime enrollment), executor credentials (authentication and a Project's credential state through `sessions.ExecutorCredentialReader`, and issuance, rotation and revocation; the Core-key Project operations record their audit entry in the same transaction) and native installation authorization and claims, whose tokens `sessionpg.Store` signs with the credential key. `cmd/server` builds one `sessionpg.Store` with the credential key and wires the Service and the Store into `execution.Dispatcher.Sessions` and `SessionsReader`, into the api fields, into Runtime enrollment and into the daemon gateway; the Worker stages Artifacts through the Service; `cmd/environment-key` builds one without the key for its credential commands. Function calls and their application receipts, the Turn execution journal, Environment initialization, connection observations and their reconciliation, Session device binding, and file-write reservation and settlement run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. ## Request handling diff --git a/services/core/cmd/server/http_routes_test.go b/services/core/cmd/server/http_routes_test.go index 6c2bed018..2af0668e5 100644 --- a/services/core/cmd/server/http_routes_test.go +++ b/services/core/cmd/server/http_routes_test.go @@ -16,6 +16,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/api" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/google/uuid" @@ -113,7 +114,7 @@ func daemonComposition(t testing.TB) http.Handler { SessionArchive: struct{ api.SessionArchive }{}, Workspaces: struct{ api.EnvironmentWorkspaces }{}, }, - Sandboxes: &api.Sandboxes{Deployment: struct{ api.Deployment }{}, NodeAllocations: struct{ api.NodeAllocations }{}, DeploymentChanges: struct{ api.DeploymentChanges }{}, + Sandboxes: &api.Sandboxes{Deployment: struct{ api.Deployment }{}, NodeAllocations: unusedNodeAllocations{}, DeploymentChanges: struct{ api.DeploymentChanges }{}, DeploymentReset: struct{ api.DeploymentReset }{}, ConfigurationDiscovery: struct{ api.ConfigurationDiscovery }{}}, }) if err != nil { @@ -314,3 +315,12 @@ func oracleTarget(segments []string, trailing bool) string { } return target.String() } + +// unusedNodeAllocations stands in for the node allocation list, which the +// route tests never read. The interface's method shares its name, so the +// embedded-interface stand-in the other dependencies use cannot satisfy it. +type unusedNodeAllocations struct{} + +func (unusedNodeAllocations) NodeAllocations(context.Context, string) ([]deployment.NodeAllocation, error) { + panic("unexpected NodeAllocations") +} diff --git a/services/core/cmd/server/installation.go b/services/core/cmd/server/installation.go index 240e7e8d3..a929acd4e 100644 --- a/services/core/cmd/server/installation.go +++ b/services/core/cmd/server/installation.go @@ -7,7 +7,7 @@ import ( "regexp" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/api" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" ) var sourceCommit = regexp.MustCompile(`^[0-9a-f]{40}$`) @@ -22,7 +22,7 @@ func installationFacts(publicURL string) (api.Installation, error) { } if publicURL != "" { base := publicURL + "/v1" - facts.PublicURL, facts.APIBaseURL, facts.LocalOnly = &publicURL, &base, deployment.LoopbackOrigin(publicURL) + facts.PublicURL, facts.APIBaseURL, facts.LocalOnly = &publicURL, &base, placement.LoopbackOrigin(publicURL) } if sourceCommit.MatchString(buildRevision) { revision := buildRevision diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index ce3c51a3a..8bf6d7c15 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/coremetrics" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/databaseurl" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmenttemplates" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/files" @@ -125,7 +126,6 @@ func run() error { return err } executionStore := store.NewWithCredentialCipher(pool, credentialKey) - executionStore.SetPublicURL(public) units := pgunit.NewPool(pool) auditStore := auditpg.New(units) agentStore := agentpg.New(units, credentialKey) @@ -159,8 +159,17 @@ func run() error { return err } sandboxProviders := providers.Builtin() + // The placement rules are built once: the provider declarations and the + // public URL never change while Core runs. + placementRules, err := placement.NewRules(sandboxProviders, public) + if err != nil { + return err + } + // Session creation in store still decides admission and placement until + // it moves to sessions. + executionStore.SetPlacement(placementRules) deploymentStore := deploymentpg.New(units, credentialKey) - deploymentService, err := deployment.NewService(deploymentStore, deploymentStore, sandboxProviders, public) + deploymentService, err := deployment.NewService(deploymentStore, deploymentStore, sandboxProviders, placementRules) if err != nil { return err } @@ -188,7 +197,7 @@ func run() error { defer func() { cancelAuditCleanup(); <-auditCleanupDone }() var workerDone chan error var worker *execution.Worker - managedNodes, err := configureManagedNodes(executionStore, deploymentService, deploymentStore, sandboxProviders, public, func(ctx context.Context) error { + managedNodes, err := configureManagedNodes(deploymentService, deploymentStore, sandboxProviders, public, func(ctx context.Context) error { if worker == nil { return errors.New("sandbox execution owner is unavailable") } @@ -204,7 +213,7 @@ func run() error { managed = managedNodes.runtime observationSources[managed.InstallationID] = managedNodes.setup } - observationResolver, err := observationstoreresolver.NewResolver(executionStore) + observationResolver, err := observationstoreresolver.NewResolver(executionStore, deploymentStore) if err != nil { return err } @@ -409,7 +418,7 @@ func run() error { if managedNodes != nil { deps.Sandboxes = &api.Sandboxes{ Deployment: deploymentService, - NodeAllocations: executionStore, + NodeAllocations: deploymentStore, DeploymentChanges: worker, DeploymentReset: worker, ConfigurationDiscovery: managedNodes.setup, diff --git a/services/core/cmd/server/managed_generations.go b/services/core/cmd/server/managed_generations.go index 3d744c15a..3a349f6ed 100644 --- a/services/core/cmd/server/managed_generations.go +++ b/services/core/cmd/server/managed_generations.go @@ -201,7 +201,7 @@ func (s *managedSetup) routeGenerations(candidate execution.PreparedRuntimeDeplo } } for after := ""; ; { - page, err := s.allocations.SandboxCredentialAllocationPage(ctx, after) + page, err := s.allocations.CredentialAllocations(ctx, after) if err != nil { return err } diff --git a/services/core/cmd/server/managed_generations_test.go b/services/core/cmd/server/managed_generations_test.go index adc1a9519..7e5855f6b 100644 --- a/services/core/cmd/server/managed_generations_test.go +++ b/services/core/cmd/server/managed_generations_test.go @@ -14,7 +14,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/e2b" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -151,7 +150,7 @@ print(json.dumps(result)) committed := current setups := &fakeDeploymentSetups{t: t, setup: committedSetup(&committed), withCredential: credentialService(t), generationPage: func(context.Context, int64) ([]deployment.Setup, error) { return nil, nil }} - allocations := &fakeGenerationAllocations{t: t, sandboxCredentialAllocationPage: func(context.Context, string) ([]store.RuntimeAllocation, error) { return nil, nil }} + allocations := &fakeGenerationAllocations{t: t, credentialAllocations: func(context.Context, string) ([]deployment.Allocation, error) { return nil, nil }} s := &managedSetup{processPaths: paths, registry: providers.Builtin(), installationID: id, deployment: setups, allocations: allocations} loaded, err := s.load(t.Context()) if err != nil { diff --git a/services/core/cmd/server/managed_nodes.go b/services/core/cmd/server/managed_nodes.go index 4ee963ac1..dc9a667ff 100644 --- a/services/core/cmd/server/managed_nodes.go +++ b/services/core/cmd/server/managed_nodes.go @@ -13,7 +13,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/node" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -26,10 +25,10 @@ type managedNodes struct { } // configureManagedNodes serves the nodes of the Web-managed deployment. Node -// presence and health go through the deployment service, the owner epoch that -// fences connections is read from the deployment reader, and the store -// resolves the generation of each allocation. -func configureManagedNodes(s *store.Store, nodes *deployment.Service, reader deployment.Reader, registry *providers.Registry, publicURL string, owner func(context.Context) error) (*managedNodes, error) { +// presence and health and the generation of each allocation go through the +// deployment service; the owner epoch that fences connections and the +// allocations each generation retains are read from the deployment reader. +func configureManagedNodes(nodes *deployment.Service, reader deployment.Reader, registry *providers.Registry, publicURL string, owner func(context.Context) error) (*managedNodes, error) { setupID := os.Getenv("OAC_INSTALLATION_ID") if setupID == "" { return nil, nil @@ -101,7 +100,7 @@ func configureManagedNodes(s *store.Store, nodes *deployment.Service, reader dep return nodes.Heartbeat(ctx, n.NodeID, connection, epoch, nodeHealthRecord(health)) }, }) - result.setup = &managedSetup{processPaths: providerProcessPaths(), registry: registry, deployment: nodes, allocations: s, hub: result.hub, installationID: setupID, publicURL: publicURL} + result.setup = &managedSetup{processPaths: providerProcessPaths(), registry: registry, deployment: nodes, allocations: reader, hub: result.hub, installationID: setupID, publicURL: publicURL} result.runtime = execution.NewDeferredRuntimeProvider(setupID, result.setup.load, result.setup.prepare) result.runtime.PublishUnconfigured = result.setup.publishUnconfigured success = true diff --git a/services/core/cmd/server/managed_setup.go b/services/core/cmd/server/managed_setup.go index 295075307..dfce8cdc6 100644 --- a/services/core/cmd/server/managed_setup.go +++ b/services/core/cmd/server/managed_setup.go @@ -16,7 +16,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/node" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) // managedSetup publishes one immutable selection to execution, bootstrap and @@ -43,13 +42,15 @@ type deploymentSetups interface { AllocationSetup(context.Context, sandbox.Reference) (deployment.Setup, error) GenerationPage(context.Context, int64) ([]deployment.Setup, error) WithCredential(owner, candidate deployment.Setup) (deployment.Setup, error) + // AllocationGeneration resolves the node and generation that own an + // allocation. + AllocationGeneration(context.Context, sandbox.Reference) (string, uint64, error) } -// generationAllocations resolves the generation that owns each allocation. -// *store.Store implements it. +// generationAllocations pages the unreleased allocations whose generations +// still hold a credential. deployment.Reader implements it. type generationAllocations interface { - ResolveRuntimeGeneration(context.Context, sandbox.Reference) (string, uint64, error) - SandboxCredentialAllocationPage(context.Context, string) ([]store.RuntimeAllocation, error) + CredentialAllocations(context.Context, string) ([]deployment.Allocation, error) } // DiscoverConfiguration asks a Provider which configuration values its @@ -178,7 +179,7 @@ func (s *managedSetup) provider(setup deployment.Setup) (sandbox.SandboxProvider if s.hub == nil { return nil, errors.New("sandbox node transport is unavailable") } - return s.hub.GenerationProvider(setup.Provider, setup.Operations, s.allocations.ResolveRuntimeGeneration), nil + return s.hub.GenerationProvider(setup.Provider, setup.Operations, s.deployment.AllocationGeneration), nil } return s.registry.BuildDirect(providers.DirectConfig{ProcessPaths: s.processPaths, InstallationID: setup.InstallationID, Selection: sandbox.Selection{Provider: setup.Provider, DeploymentSpec: setup.Specification, Configuration: setup.Configuration}, Fence: &s.providerCalls}) } diff --git a/services/core/cmd/server/managed_setup_test.go b/services/core/cmd/server/managed_setup_test.go index 1b779186b..4f117624e 100644 --- a/services/core/cmd/server/managed_setup_test.go +++ b/services/core/cmd/server/managed_setup_test.go @@ -11,6 +11,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/providercontract" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" @@ -22,7 +23,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/node" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -34,10 +34,10 @@ func TestWebSetupCreatesManagerWithoutLocalProvider(t *testing.T) { t.Fatal(err) } t.Setenv("OAC_CORE_KEY_DIGESTS_FILE", path) - if _, err := configureManagedNodes(nil, nil, nil, providers.Builtin(), "", nil); err == nil || !strings.Contains(err.Error(), "OAC_PUBLIC_URL") { + if _, err := configureManagedNodes(nil, nil, providers.Builtin(), "", nil); err == nil || !strings.Contains(err.Error(), "OAC_PUBLIC_URL") { t.Fatal("sandbox manager started without a public URL", err) } - m, err := configureManagedNodes(nil, nil, nil, providers.Builtin(), "https://core.example", func(context.Context) error { return nil }) + m, err := configureManagedNodes(nil, nil, providers.Builtin(), "https://core.example", func(context.Context) error { return nil }) if err != nil { t.Fatal(err) } @@ -46,7 +46,7 @@ func TestWebSetupCreatesManagerWithoutLocalProvider(t *testing.T) { t.Fatal("zero-node setup unexpectedly instantiated local compute or omitted management") } t.Setenv("OAC_CORE_KEY_DIGESTS_FILE", "") - if _, err := configureManagedNodes(nil, nil, nil, providers.Builtin(), "https://core.example", nil); err == nil { + if _, err := configureManagedNodes(nil, nil, providers.Builtin(), "https://core.example", nil); err == nil { t.Fatal("setup accepted without admin authentication") } } @@ -54,11 +54,12 @@ func TestWebSetupCreatesManagerWithoutLocalProvider(t *testing.T) { // fakeDeploymentSetups is a strict deploymentSetups: a call without a set // function fails the test. type fakeDeploymentSetups struct { - t testing.TB - setup func(context.Context) (deployment.Setup, error) - allocationSetup func(context.Context, sandbox.Reference) (deployment.Setup, error) - generationPage func(context.Context, int64) ([]deployment.Setup, error) - withCredential func(owner, candidate deployment.Setup) (deployment.Setup, error) + t testing.TB + setup func(context.Context) (deployment.Setup, error) + allocationSetup func(context.Context, sandbox.Reference) (deployment.Setup, error) + generationPage func(context.Context, int64) ([]deployment.Setup, error) + withCredential func(owner, candidate deployment.Setup) (deployment.Setup, error) + allocationGeneration func(context.Context, sandbox.Reference) (string, uint64, error) } func (f *fakeDeploymentSetups) Setup(ctx context.Context) (deployment.Setup, error) { @@ -85,25 +86,24 @@ func (f *fakeDeploymentSetups) WithCredential(owner, candidate deployment.Setup) } return f.withCredential(owner, candidate) } +func (f *fakeDeploymentSetups) AllocationGeneration(ctx context.Context, ref sandbox.Reference) (string, uint64, error) { + if f.allocationGeneration == nil { + return "", 0, unexpectedCall(f.t, "AllocationGeneration") + } + return f.allocationGeneration(ctx, ref) +} // fakeGenerationAllocations is a strict generationAllocations. type fakeGenerationAllocations struct { - t testing.TB - resolveRuntimeGeneration func(context.Context, sandbox.Reference) (string, uint64, error) - sandboxCredentialAllocationPage func(context.Context, string) ([]store.RuntimeAllocation, error) + t testing.TB + credentialAllocations func(context.Context, string) ([]deployment.Allocation, error) } -func (f *fakeGenerationAllocations) ResolveRuntimeGeneration(ctx context.Context, ref sandbox.Reference) (string, uint64, error) { - if f.resolveRuntimeGeneration == nil { - return "", 0, unexpectedCall(f.t, "ResolveRuntimeGeneration") - } - return f.resolveRuntimeGeneration(ctx, ref) -} -func (f *fakeGenerationAllocations) SandboxCredentialAllocationPage(ctx context.Context, after string) ([]store.RuntimeAllocation, error) { - if f.sandboxCredentialAllocationPage == nil { - return nil, unexpectedCall(f.t, "SandboxCredentialAllocationPage") +func (f *fakeGenerationAllocations) CredentialAllocations(ctx context.Context, after string) ([]deployment.Allocation, error) { + if f.credentialAllocations == nil { + return nil, unexpectedCall(f.t, "CredentialAllocations") } - return f.sandboxCredentialAllocationPage(ctx, after) + return f.credentialAllocations(ctx, after) } // unexpectedCall fails the test from any goroutine and returns the error the @@ -122,7 +122,11 @@ func committedSetup(value *deployment.Setup) func(context.Context) (deployment.S // deployment service does. That needs no storage. func credentialService(t *testing.T) func(owner, candidate deployment.Setup) (deployment.Setup, error) { t.Helper() - service, err := deployment.NewService(deploymentpg.New(nil, nil), deploymentpg.New(nil, nil), providers.Builtin(), "") + rules, err := placement.NewRules(providers.Builtin(), "") + if err != nil { + t.Fatal(err) + } + service, err := deployment.NewService(deploymentpg.New(nil, nil), deploymentpg.New(nil, nil), providers.Builtin(), rules) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/api/core_store_validation_test.go b/services/core/internal/api/core_store_validation_test.go index 0cca0fb45..2279170a5 100644 --- a/services/core/internal/api/core_store_validation_test.go +++ b/services/core/internal/api/core_store_validation_test.go @@ -11,6 +11,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" @@ -20,7 +21,11 @@ import ( func TestCoreStoreValidationFieldsAndPublicFallback(t *testing.T) { // The deployment validates these inputs before it reaches storage, so // storage without a database is enough. - nodes, err := deployment.NewService(deploymentpg.New(nil, nil), deploymentpg.New(nil, nil), providers.Builtin(), "") + rules, err := placement.NewRules(providers.Builtin(), "") + if err != nil { + t.Fatal(err) + } + nodes, err := deployment.NewService(deploymentpg.New(nil, nil), deploymentpg.New(nil, nil), providers.Builtin(), rules) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/api/errors.go b/services/core/internal/api/errors.go index 13632f42b..9b34108a5 100644 --- a/services/core/internal/api/errors.go +++ b/services/core/internal/api/errors.go @@ -9,6 +9,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/obs/log" "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/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/textvalue" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/writeaudit" @@ -89,10 +90,14 @@ func writeFieldError(w http.ResponseWriter, err error) bool { // reset or the sandbox deployment, then the Session errors. func writeStoreError(w http.ResponseWriter, r *http.Request, err error) { // Admission paused by a reset is Session admission's. - if errors.Is(err, store.ErrSandboxResetAdmission) { + if errors.Is(err, placement.ErrResetAdmission) { writeError(w, http.StatusServiceUnavailable, "sandbox_reset_in_progress", "A sandbox reset is in progress.") return } + if errors.Is(err, placement.ErrAdmissionClosed) { + writeError(w, http.StatusConflict, "environment_unavailable", "The environment is no longer available for new input.") + return + } if writeSandboxError(w, err) { return } diff --git a/services/core/internal/api/errors_deployment.go b/services/core/internal/api/errors_deployment.go index 873c9879c..352b39a2b 100644 --- a/services/core/internal/api/errors_deployment.go +++ b/services/core/internal/api/errors_deployment.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/providercontract" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" @@ -85,7 +86,7 @@ func writeSandboxError(w http.ResponseWriter, err error) bool { writeError(w, http.StatusConflict, "sandbox_not_configured", "The sandbox deployment is not configured.") case errors.As(err, &deploymentConfiguration): writeError(w, http.StatusBadRequest, "invalid_sandbox_configuration", deploymentConfiguration.Message) - case errors.Is(err, deployment.ErrPublicURLUnreachable): + case errors.Is(err, placement.ErrPublicURLUnreachable): writeError(w, http.StatusConflict, "sandbox_configuration_error", err.Error()) case errors.Is(err, deployment.ErrNodeAddressMismatch): writeError(w, http.StatusConflict, "sandbox_node_address_mismatch", "This node uses a different Core address than the installation public URL. Generate a new command on the Nodes page and run it on the host.") @@ -101,9 +102,9 @@ func writeSandboxError(w http.ResponseWriter, err error) bool { writeError(w, http.StatusConflict, "runtime_node_in_use", "The sandbox node retains allocations, snapshots, reservations or pending cleanup.") case errors.Is(err, deployment.ErrLocalNodeConfigured): writeError(w, http.StatusConflict, "runtime_local_node_configured", "The local sandbox node is enabled in deployment configuration. Drain it with the previous release and remove its file-managed configuration before replacing it.") - case errors.Is(err, deployment.ErrNodesPreparing): + case errors.Is(err, placement.ErrNodesPreparing): writeError(w, http.StatusServiceUnavailable, "sandbox_nodes_preparing", "Sandbox nodes are preparing the requested Runtime.") - case errors.Is(err, deployment.ErrNodeUnavailable): + case errors.Is(err, placement.ErrNodeUnavailable): writeError(w, http.StatusServiceUnavailable, "runtime_node_unavailable", "The selected sandbox node is unavailable or has no capacity.") default: return false diff --git a/services/core/internal/api/fakes_test.go b/services/core/internal/api/fakes_test.go index a4fd14664..e51df824c 100644 --- a/services/core/internal/api/fakes_test.go +++ b/services/core/internal/api/fakes_test.go @@ -266,15 +266,15 @@ func (f *fakeDeployment) DecodeConfiguration(a0 string, a1 json.RawMessage, a2 j } type fakeNodeAllocations struct { - t testing.TB - listNodeRuntimeAllocations func(context.Context, string) ([]store.RuntimeNodeAllocation, error) + t testing.TB + nodeAllocations func(context.Context, string) ([]deployment.NodeAllocation, error) } -func (f *fakeNodeAllocations) ListNodeRuntimeAllocations(a0 context.Context, a1 string) ([]store.RuntimeNodeAllocation, error) { - if f.listNodeRuntimeAllocations == nil { - unexpectedCall(f.t, "ListNodeRuntimeAllocations") +func (f *fakeNodeAllocations) NodeAllocations(a0 context.Context, a1 string) ([]deployment.NodeAllocation, error) { + if f.nodeAllocations == nil { + unexpectedCall(f.t, "NodeAllocations") } - return f.listNodeRuntimeAllocations(a0, a1) + return f.nodeAllocations(a0, a1) } type fakeDeploymentChanges struct { diff --git a/services/core/internal/api/installation_test.go b/services/core/internal/api/installation_test.go index 74dd81a6b..94cd2bf7c 100644 --- a/services/core/internal/api/installation_test.go +++ b/services/core/internal/api/installation_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" ) @@ -70,7 +71,7 @@ func TestDeploymentAddressIsNotInput(t *testing.T) { if input.Provider != "e2b" || input.ExpectedGeneration != 0 { t.Fatal("invalid selection reached initialization", input.Provider, input.ExpectedGeneration) } - return deployment.View{}, deployment.ErrPublicURLUnreachable + return deployment.View{}, placement.ErrPublicURLUnreachable } // The strict fake fails the test if core_url reaches the update. fakes.deploymentChanges.initializeSandboxDeployment = initialize diff --git a/services/core/internal/api/placement_errors_test.go b/services/core/internal/api/placement_errors_test.go new file mode 100644 index 000000000..36866dae5 --- /dev/null +++ b/services/core/internal/api/placement_errors_test.go @@ -0,0 +1,48 @@ +package api + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" +) + +// Each placement error keeps its status, code and message, wrapped as the +// rules and their callers return it. Session creation reports them through +// writeStoreError; the deployment writer reports those deployment use cases +// can meet. +func TestPlacementErrorsKeepTheirResponses(t *testing.T) { + for _, tc := range []struct { + err error + status int + code, message string + deploymentWrite bool + }{ + {placement.ErrResetAdmission, http.StatusServiceUnavailable, "sandbox_reset_in_progress", "A sandbox reset is in progress.", false}, + {fmt.Errorf("%w: sandbox creation is paused for provider maintenance", placement.ErrAdmissionClosed), http.StatusConflict, "environment_unavailable", "The environment is no longer available for new input.", false}, + {placement.ErrPublicURLUnreachable, http.StatusConflict, "sandbox_configuration_error", placement.ErrPublicURLUnreachable.Error(), true}, + {placement.ErrNodesPreparing, http.StatusServiceUnavailable, "sandbox_nodes_preparing", "Sandbox nodes are preparing the requested Runtime.", true}, + {placement.ErrNodeUnavailable, http.StatusServiceUnavailable, "runtime_node_unavailable", "The selected sandbox node is unavailable or has no capacity.", true}, + } { + writers := map[string]func(http.ResponseWriter, *http.Request, error){"store": writeStoreError} + if tc.deploymentWrite { + writers["deployment"] = writeDeploymentError + } + for name, write := range writers { + response := httptest.NewRecorder() + write(response, httptest.NewRequest(http.MethodPost, "/v1/agents/sessions", nil), tc.err) + var body struct { + Error struct{ Code, Message string } `json:"error"` + } + if err := json.Unmarshal(response.Body.Bytes(), &body); err != nil { + t.Fatal(name, err) + } + if response.Code != tc.status || body.Error.Code != tc.code || body.Error.Message != tc.message { + t.Errorf("%s writer for %v = %d %+v", name, tc.err, response.Code, body.Error) + } + } + } +} diff --git a/services/core/internal/api/sandbox_deployment_changes_test.go b/services/core/internal/api/sandbox_deployment_changes_test.go index c83424794..e7b7d4ed2 100644 --- a/services/core/internal/api/sandbox_deployment_changes_test.go +++ b/services/core/internal/api/sandbox_deployment_changes_test.go @@ -8,10 +8,10 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/e2b" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestSandboxDeploymentChangesAuthenticateAndDecode(t *testing.T) { @@ -89,7 +89,7 @@ func TestSandboxMutationErrorsExposeOnlyTypedCoreFacts(t *testing.T) { {&deployment.ResetRequiredError{CurrentProvider: "e2b", RequestedProvider: "e2b"}, "sandbox_reset_required", 409, `"current_provider":"e2b"`, false}, {&deployment.InUseError{Resources: deployment.Resources{Allocations: 2, Pending: 1}}, "sandbox_in_use", 409, `"allocations":2`, false}, {deployment.ErrResetInProgress, "sandbox_reset_in_progress", 409, "", false}, - {store.ErrSandboxResetAdmission, "sandbox_reset_in_progress", 503, "", true}, + {placement.ErrResetAdmission, "sandbox_reset_in_progress", 503, "", true}, } { for name, write := range writers { if tc.admission && name != "store" { diff --git a/services/core/internal/api/sandbox_manager.go b/services/core/internal/api/sandbox_manager.go index 4d008d197..99c814d47 100644 --- a/services/core/internal/api/sandbox_manager.go +++ b/services/core/internal/api/sandbox_manager.go @@ -8,7 +8,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/go-chi/chi/v5" ) @@ -16,7 +15,7 @@ type SandboxNodeList struct { Data []deployment.Node `json:"data"` } type SandboxAllocationList struct { - Data []store.RuntimeNodeAllocation `json:"data"` + Data []deployment.NodeAllocation `json:"data"` } type SandboxEnrollmentToken struct { Token string `json:"token"` @@ -54,7 +53,7 @@ type Deployment interface { // NodeAllocations lists the allocations a sandbox node holds. type NodeAllocations interface { - ListNodeRuntimeAllocations(ctx context.Context, nodeID string) ([]store.RuntimeNodeAllocation, error) + NodeAllocations(ctx context.Context, nodeID string) ([]deployment.NodeAllocation, error) } // registerSandboxNodeRoutes serves node machine connections. They authenticate @@ -179,9 +178,9 @@ func (h *Handler) removeSandboxNode(w http.ResponseWriter, r *http.Request) { // @Failure 400,401,404,409,500,503 {object} CoreErrorResponse // @Router /core/v1/sandbox/nodes/{node_id}/allocations [get] func (h *Handler) sandboxAllocations(w http.ResponseWriter, r *http.Request) { - value, err := h.Sandboxes.NodeAllocations.ListNodeRuntimeAllocations(r.Context(), chi.URLParam(r, "node_id")) + value, err := h.Sandboxes.NodeAllocations.NodeAllocations(r.Context(), chi.URLParam(r, "node_id")) if err != nil { - writeStoreError(w, r, err) + writeDeploymentError(w, r, err) return } writeJSON(w, http.StatusOK, SandboxAllocationList{Data: value}) diff --git a/services/core/internal/db/queries/runtime_allocations.sql b/services/core/internal/db/queries/runtime_allocations.sql index 0296b6249..fc62d9deb 100644 --- a/services/core/internal/db/queries/runtime_allocations.sql +++ b/services/core/internal/db/queries/runtime_allocations.sql @@ -52,11 +52,3 @@ WHERE id = $1 AND state <> 'released' RETURNING *; -- name: ReleaseRuntimeAllocation :one UPDATE runtime_allocations SET state = 'released', released_at = clock_timestamp() WHERE id = $1 AND state = 'cleanup_pending' AND create_settled RETURNING *; - --- name: ListUnallocatedHostedEnvironments :many -SELECT e.id, s.tenant_id -FROM environments e JOIN sessions s ON s.id = e.session_id -WHERE e.id > $1 AND s.deleted_at IS NULL AND e.status = 'pending' - AND s.configuration->'environment'->>'type' = 'openai_hosted' - AND NOT EXISTS (SELECT 1 FROM runtime_allocations a WHERE a.environment_id = e.id) -ORDER BY e.id LIMIT 32; diff --git a/services/core/internal/db/sqlc/runtime_allocations.sql.go b/services/core/internal/db/sqlc/runtime_allocations.sql.go index 7e47aef04..0caae4420 100644 --- a/services/core/internal/db/sqlc/runtime_allocations.sql.go +++ b/services/core/internal/db/sqlc/runtime_allocations.sql.go @@ -249,40 +249,6 @@ func (q *Queries) ListRuntimeObservationSessions(ctx context.Context, arg ListRu return items, nil } -const listUnallocatedHostedEnvironments = `-- name: ListUnallocatedHostedEnvironments :many -SELECT e.id, s.tenant_id -FROM environments e JOIN sessions s ON s.id = e.session_id -WHERE e.id > $1 AND s.deleted_at IS NULL AND e.status = 'pending' - AND s.configuration->'environment'->>'type' = 'openai_hosted' - AND NOT EXISTS (SELECT 1 FROM runtime_allocations a WHERE a.environment_id = e.id) -ORDER BY e.id LIMIT 32 -` - -type ListUnallocatedHostedEnvironmentsRow struct { - ID pgtype.UUID `json:"id"` - TenantID pgtype.UUID `json:"tenant_id"` -} - -func (q *Queries) ListUnallocatedHostedEnvironments(ctx context.Context, id pgtype.UUID) ([]ListUnallocatedHostedEnvironmentsRow, error) { - rows, err := q.db.Query(ctx, listUnallocatedHostedEnvironments, id) - if err != nil { - return nil, err - } - defer rows.Close() - items := []ListUnallocatedHostedEnvironmentsRow{} - for rows.Next() { - var i ListUnallocatedHostedEnvironmentsRow - if err := rows.Scan(&i.ID, &i.TenantID); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - const observeRuntimeRunning = `-- name: ObserveRuntimeRunning :one UPDATE runtime_allocations SET state = 'running', create_settled = true WHERE id = $1 AND state IN ('creating', 'running') diff --git a/services/core/internal/store/runtime_idle_policy_test.go b/services/core/internal/deployment/activity_test.go similarity index 88% rename from services/core/internal/store/runtime_idle_policy_test.go rename to services/core/internal/deployment/activity_test.go index b08b740cf..2f308dfb2 100644 --- a/services/core/internal/store/runtime_idle_policy_test.go +++ b/services/core/internal/deployment/activity_test.go @@ -1,4 +1,4 @@ -package store +package deployment import ( "testing" @@ -28,7 +28,7 @@ func TestRuntimeIdleAdmissionUsesDatabaseClock(t *testing.T) { {"invalid_policy", timeout + time.Second, false, false, true, -time.Second, false}, } { t.Run(test.name, func(t *testing.T) { - activity := RuntimeActivity{ObservedAt: observed, LastActivity: observed.Add(-test.age), Busy: test.busy, WakeRequested: test.wake, HasCompletedTurn: test.completed} + activity := Activity{ObservedAt: observed, LastActivity: observed.Add(-test.age), Busy: test.busy, WakeRequested: test.wake, HasCompletedTurn: test.completed} if got := activity.ReadyToSuspend(test.timeout); got != test.want { t.Fatalf("suspension admission = %v, want %v with database clock skew %s", got, test.want, skew) } diff --git a/services/core/internal/deployment/allocation.go b/services/core/internal/deployment/allocation.go new file mode 100644 index 000000000..c825dda4f --- /dev/null +++ b/services/core/internal/deployment/allocation.go @@ -0,0 +1,175 @@ +package deployment + +import ( + "encoding/json" + "time" +) + +// Allocation retains a hosted Environment's compute ownership, not its public +// readiness. It survives Session deletion until cleanup is confirmed and keeps +// no bootstrap secret. +type Allocation struct { + DeploymentGeneration uint64 + // NodeID is empty for an allocation no node serves. + NodeID string + ObservationError string + ComputePhase string + ComputeRevision int64 + ComputeState json.RawMessage + ComputeActivityAt time.Time + ComputeWakeRequested bool + ComputeRetainedUntil *time.Time + ID, EnvironmentID string + SessionID, TenantID string + DeviceID string + // ProviderKey is the installation the allocation was provisioned for. + ProviderKey string + State string + // CreateSettled records evidence that the original Create can no longer + // change resources. + CreateSettled bool + // SessionDeleted reports that the owning Session was publicly deleted. + SessionDeleted bool + // Replayed reports that ReserveAllocation returned an existing allocation, + // which never authorizes another Create. + Replayed bool + // Expired reports, by the database clock, that the allocation's lease or + // retention has passed. + Expired bool + CreatedAt, KeptAt time.Time +} + +// Key names the allocation's Environment. +func (a Allocation) Key() AllocationKey { + return AllocationKey{TenantID: a.TenantID, EnvironmentID: a.EnvironmentID} +} + +// AllocationKey names a tenant's hosted Environment, which has at most one +// allocation. +type AllocationKey struct{ TenantID, EnvironmentID string } + +// parse returns the key in canonical form, and ErrInvalidInput for a +// malformed identifier. +func (k AllocationKey) parse() (AllocationKey, error) { + tenant, err := parseID(k.TenantID) + if err != nil { + return AllocationKey{}, err + } + environment, err := parseID(k.EnvironmentID) + if err != nil { + return AllocationKey{}, err + } + return AllocationKey{TenantID: tenant, EnvironmentID: environment}, nil +} + +// NewAllocation is an allocation to store with its dedicated device. +type NewAllocation struct { + ID, EnvironmentID, DeviceID, ProviderKey string + // NodeID is empty when no node serves the allocation. + NodeID string + Generation uint64 +} + +// SessionDevice is the Runtime device a Session is bound to. +type SessionDevice struct{ ID, EnvironmentID string } + +// Activity separates real work from the connection keepalive. ObservedAt is +// the database clock when it was read. +type Activity struct { + LastActivity, ObservedAt time.Time + Busy, WakeRequested, HasCompletedTurn bool +} + +// ReadyToSuspend evaluates the idle policy on one database-clock observation. +func (a Activity) ReadyToSuspend(idleTimeout time.Duration) bool { + return idleTimeout > 0 && a.HasCompletedTurn && !a.Busy && !a.WakeRequested && + a.ObservedAt.Sub(a.LastActivity) >= idleTimeout +} + +// ComputeChange is the compute phase an operation commits before its external +// effects. +type ComputeChange struct { + Phase string + State json.RawMessage + // RetainedUntil is nil only for the running phase. + RetainedUntil *time.Time +} + +// ObservationSession is the minimum durable identity the deployment-wide +// read-only sampler needs. The observation service resolves the provider +// identity again before any external read. +type ObservationSession struct{ TenantID, SessionID string } + +// ObservationSessionPage is one page of ObservationSessions. NextCursor is +// empty on the last page. +type ObservationSessionPage struct { + Sessions []ObservationSession + NextCursor string +} + +// UnallocatedEnvironment is a committed hosted Environment awaiting its +// allocation. A missing allocation differs from an unknown Create outcome of +// an existing one. +type UnallocatedEnvironment struct{ ID, TenantID string } + +// LifecyclePlacement is what routing a hosted Environment to its node +// lifecycle reads: the deployment's provider and mode, the Environment's +// allocation and the node placement its Session reserved. +type LifecyclePlacement struct { + Provider, Mode string + // AllocationID and AllocationNodeID are empty without an allocation or + // its node. + AllocationID, AllocationNodeID string + // PlacementNodeID is empty without a placement. + PlacementNodeID string + PlacementReleased bool +} + +// NodeAllocation is an unreleased allocation a node serves. +type NodeAllocation struct { + DeploymentGeneration uint64 `json:"deployment_generation"` + Diagnostic string `json:"diagnostic"` + ID string `json:"id"` + NodeID string `json:"node_id"` + TenantID string `json:"tenant_id"` + SessionID string `json:"session_id"` + EnvironmentID string `json:"environment_id"` + State string `json:"state"` + ComputePhase string `json:"compute_phase"` + // The time the allocation entered its current compute_phase, or null when unknown; an allocation that existed before Core recorded it reports null until its next phase change. For a suspended microsandbox allocation, this time plus the deployment's snapshot retention tells roughly when Core reclaims it. + ComputePhaseChangedAt *time.Time `json:"compute_phase_changed_at" extensions:"x-nullable"` + Initialization string `json:"initialization"` + CreatedAt time.Time `json:"created_at"` +} + +// observationDiagnostics are the diagnostics an observation records; empty +// clears the previous one. +var observationDiagnostics = map[string]bool{ + "": true, "node_unavailable": true, "resource_missing": true, + "compute_unconfirmed": true, "ownership_mismatch": true, "provider_unavailable": true, +} + +// computeTransition reports whether managed compute may move from one phase +// to another. Repeating a phase records a new state, except while disabled. +func computeTransition(from, to string) bool { + if from == to { + return from != "disabled" + } + switch from { + case "disabled": + return to == "running" + case "running": + return to == "quiescing" + case "quiescing": + return to == "suspending" || to == "waking" || to == "running" + case "suspending": + return to == "suspended" || to == "waking" + case "suspended": + return to == "restoring" + case "restoring": + return to == "waking" + case "waking": + return to == "running" + } + return false +} diff --git a/services/core/internal/deployment/allocations.go b/services/core/internal/deployment/allocations.go new file mode 100644 index 000000000..294c096fc --- /dev/null +++ b/services/core/internal/deployment/allocations.go @@ -0,0 +1,375 @@ +package deployment + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/google/uuid" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// ReserveAllocation commits the allocation of the tenant's hosted Environment +// and its dedicated device together, before the provider creates compute. +// installation is the provider key the allocation is provisioned for, and +// credentialHash the SHA-256 digest of the device credential. An existing +// allocation for the same installation returns Replayed; only a fresh one +// authorizes the one Create call. +func (e *ExecutionOperations) ReserveAllocation(ctx context.Context, key AllocationKey, installation, credentialHash string) (Allocation, error) { + key, err := key.parse() + if err != nil { + return Allocation{}, err + } + if installation, err = parseID(installation); err != nil { + return Allocation{}, err + } + registration, err := sessions.NewDeviceRegistration("managed-runtime", credentialHash) + if err != nil { + return Allocation{}, fmt.Errorf("%w: SHA-256 credential digest required", ErrInvalidInput) + } + var result Allocation + err = e.storage.WithReservation(ctx, key, func(locked sessions.LockedSession, tx ReservationTx) error { + if err := locked.Public(); err != nil { + return err + } + environment, err := tx.LoadEnvironment(ctx) + if err != nil { + return err + } + if kind, err := sessions.EnvironmentType(environment.Configuration); err != nil || kind != "openai_hosted" { + return ErrInvalidInput + } + existing, found, err := tx.FindAllocation() + if err != nil { + return err + } + if found { + if existing.ProviderKey != installation { + return ErrAllocationConflict + } + existing.Replayed = true + result = existing + return nil + } + // Existing receipts replay first, so retries and cleanup continue + // while admission is closed. + d, err := tx.LockDeployment() + if err != nil { + return err + } + if err := e.service.rules.CheckAdmission(d, installation); err != nil { + return err + } + if environment.Status == "failed" || environment.Status == "expired" { + return ErrInvalidInput + } + allocation := NewAllocation{ID: uuid.NewString(), EnvironmentID: environment.ID, DeviceID: uuid.NewString(), ProviderKey: installation, Generation: d.Generation} + if d.Mode == "nodes" { + reserved, err := tx.LoadReserved() + if err != nil { + return err + } + if err := placement.CheckReserved(reserved); err != nil { + return err + } + allocation.NodeID, allocation.Generation = reserved.NodeID, reserved.Generation + } + device := sessions.ExecutionDevice{ID: allocation.DeviceID, Name: registration.Name, EnvironmentID: environment.ID} + if err := sessions.CreateEnvironmentDevice(ctx, tx, device, registration.CredentialHash); err != nil { + return err + } + result, err = tx.InsertAllocation(allocation) + return err + }) + if err != nil { + return Allocation{}, err + } + return result, nil +} + +// ObserveRunning records verified original compute identity. It neither +// marks the Environment connected nor qualifies native preparation. +func (e *ExecutionOperations) ObserveRunning(ctx context.Context, owner Allocation) (Allocation, error) { + return e.change(ctx, owner, true, func(tx AllocationTx, current Allocation) (Allocation, error) { + return tx.ObserveRunning(current) + }) +} + +// KeepAllocation follows an authenticated connection and a successful +// provider observation. A keepalive never revives cleanup or an expired lease. +func (e *ExecutionOperations) KeepAllocation(ctx context.Context, owner Allocation) (Allocation, error) { + return e.change(ctx, owner, true, func(tx AllocationTx, current Allocation) (Allocation, error) { + return tx.Keep(current) + }) +} + +// SettleCreation records evidence that the original Create can no longer +// change resources. A timeout, missing compute or lost lease is not evidence. +func (e *ExecutionOperations) SettleCreation(ctx context.Context, owner Allocation) (Allocation, error) { + return e.change(ctx, owner, false, func(tx AllocationTx, current Allocation) (Allocation, error) { + if current.State == "released" { + return current, nil + } + return tx.SettleCreation(current) + }) +} + +// ReleaseAllocation follows successful cleanup of the owned compute and +// volumes. An unknown Create outcome keeps its cleanup record even when the +// compute is absent. +func (e *ExecutionOperations) ReleaseAllocation(ctx context.Context, owner Allocation) (Allocation, error) { + return e.change(ctx, owner, false, func(tx AllocationTx, current Allocation) (Allocation, error) { + if current.State == "released" { + return current, nil + } + return tx.Release(current) + }) +} + +// RequestCleanup revokes the allocation's authority before external +// reclamation and settles the owning Session: a deleted Session's work is +// cancelled, and a live Environment fails with the generic provisioning +// reason unless its compute expired. Cancellation requests do not prove that +// native work has stopped. +func (e *ExecutionOperations) RequestCleanup(ctx context.Context, owner Allocation) (Allocation, error) { + return e.cleanup(ctx, owner, false) +} + +// ReleaseAbsentCreation consumes provider proof that the original Create is +// settled and owns no resources: it requests cleanup and releases the +// allocation in one commit. +func (e *ExecutionOperations) ReleaseAbsentCreation(ctx context.Context, owner Allocation) (Allocation, error) { + return e.cleanup(ctx, owner, true) +} + +func (e *ExecutionOperations) cleanup(ctx context.Context, owner Allocation, absent bool) (Allocation, error) { + key, err := owner.Key().parse() + if err != nil { + return Allocation{}, err + } + var result Allocation + err = e.storage.WithAllocationCleanup(ctx, key, func(tx AllocationCleanupTx) error { + current, err := checkAllocation(tx, owner, false) + if err != nil { + return err + } + if current.State == "released" { + result = current + return nil + } + if err := tx.RevokeDevice(current); err != nil { + return err + } + // The revocation commits with the Session's settlement, which reads + // the allocation's expiry after it. + revoked, err := tx.LoadAllocation() + if err != nil { + return err + } + if revoked.SessionDeleted { + err = sessions.CancelWork(ctx, tx) + } else { + err = sessions.TerminateEnvironment(ctx, tx, revoked.Expired, sessions.ProvisioningFailureReason, nil) + } + if err != nil { + return err + } + pending, err := tx.RequestCleanup(current) + if err != nil || !absent { + result = pending + return err + } + settled, err := tx.SettleCreation(pending) + if err != nil { + return err + } + result, err = tx.Release(settled) + return err + }) + if err != nil { + return Allocation{}, err + } + return result, nil +} + +// SetCompute commits a compute phase before its external effects. The owner's +// revision and the Session lock fence a stale lifecycle observation. Moving +// running compute to quiescing requires the idle timeout, which the +// transaction checks again against the activity it reads. +func (e *ExecutionOperations) SetCompute(ctx context.Context, owner Allocation, phase string, state json.RawMessage, retainedUntil *time.Time, idleTimeout time.Duration) (Allocation, error) { + if !computeTransition(owner.ComputePhase, phase) || !json.Valid(state) || (owner.ComputePhase == "running" && phase == "quiescing" && idleTimeout <= 0) { + return Allocation{}, ErrInvalidInput + } + if phase != "running" && (retainedUntil == nil || retainedUntil.IsZero()) { + return Allocation{}, ErrInvalidInput + } + var object map[string]json.RawMessage + if json.Unmarshal(state, &object) != nil || object == nil { + return Allocation{}, ErrInvalidInput + } + return e.change(ctx, owner, true, func(tx AllocationTx, current Allocation) (Allocation, error) { + if current.ComputeRevision != owner.ComputeRevision || current.ComputePhase != owner.ComputePhase { + return Allocation{}, ErrAllocationConflict + } + if (phase == "quiescing" && current.ComputePhase == "running") || (phase == "suspending" && current.ComputePhase == "quiescing") { + activity, err := tx.LoadActivity(current) + if err != nil { + return Allocation{}, err + } + if activity.Busy || activity.WakeRequested { + return Allocation{}, ErrAllocationConflict + } + if phase == "quiescing" && (!activity.ReadyToSuspend(idleTimeout) || current.ComputeActivityAt.After(owner.ComputeActivityAt)) { + return Allocation{}, ErrAllocationConflict + } + } + // A node-backed restore reserves capacity on its original node. + if current.ComputePhase == "suspended" && phase == "restoring" && current.NodeID != "" { + restore, err := tx.LoadRestore(current) + if err != nil { + return Allocation{}, err + } + if err := placement.CheckRestore(restore); err != nil { + return Allocation{}, err + } + } + return tx.SetCompute(current, ComputeChange{Phase: phase, State: state, RetainedUntil: retainedUntil}) + }) +} + +// RecordObservation records a diagnostic for the observed compute revision +// and state of a node-backed allocation. It never releases ownership, changes +// public readiness or authorizes replacement. +func (e *ExecutionOperations) RecordObservation(ctx context.Context, owner Allocation, diagnostic string) error { + if !observationDiagnostics[diagnostic] { + return ErrInvalidInput + } + if owner.NodeID == "" { + return nil + } + _, err := e.change(ctx, owner, false, func(tx AllocationTx, current Allocation) (Allocation, error) { + if current.ComputeRevision != owner.ComputeRevision || current.State != owner.State || current.State == "released" { + return current, nil + } + return current, tx.RecordObservation(current, diagnostic) + }) + return err +} + +// ClearWake consumes the allocation's wake request observed with the given +// activity time. A newer request stays. +func (e *ExecutionOperations) ClearWake(ctx context.Context, owner Allocation, observed time.Time) error { + id, err := parseID(owner.ID) + if err != nil { + return err + } + return e.storage.ClearWake(ctx, id, observed) +} + +// change runs apply on the owner's allocation in its Session-locked +// transaction, after checkAllocation. +func (e *ExecutionOperations) change(ctx context.Context, owner Allocation, live bool, apply func(AllocationTx, Allocation) (Allocation, error)) (Allocation, error) { + key, err := owner.Key().parse() + if err != nil { + return Allocation{}, err + } + var result Allocation + err = e.storage.WithAllocation(ctx, key, func(tx AllocationTx) error { + current, err := checkAllocation(tx, owner, live) + if err != nil { + return err + } + result, err = apply(tx, current) + return err + }) + if err != nil { + return Allocation{}, err + } + return result, nil +} + +// checkAllocation returns the stored allocation when it is still the owner's: +// the same allocation, device, installation and node. A live change also +// requires the Session undeleted and bound to the allocation's device. +func checkAllocation(tx AllocationTx, owner Allocation, live bool) (Allocation, error) { + current, err := tx.LoadAllocation() + if err != nil { + return Allocation{}, err + } + if current.ID != owner.ID || current.DeviceID != owner.DeviceID || current.ProviderKey != owner.ProviderKey || current.NodeID != owner.NodeID { + return Allocation{}, ErrAllocationConflict + } + if !live { + return current, nil + } + if current.SessionDeleted { + // sessions.ErrNotFound keeps the public 404 that writeStoreError maps. + return Allocation{}, sessions.ErrNotFound + } + device, bound, err := tx.LoadSessionDevice() + if err != nil { + return Allocation{}, err + } + if !bound || device.ID != current.DeviceID || device.EnvironmentID != current.EnvironmentID { + return Allocation{}, ErrAllocationConflict + } + return current, nil +} + +// TouchActivity records live-compute activity of the tenant's Environment. +// Only operations that need live compute call it; history and published +// Artifact reads never do. +func (s *Service) TouchActivity(ctx context.Context, tenant, environment string) error { + key, err := AllocationKey{TenantID: tenant, EnvironmentID: environment}.parse() + if err != nil { + return err + } + return s.storage.WithActivity(ctx, key, func(locked sessions.LockedSession, tx ActivityTx) error { + if err := locked.Public(); err != nil { + return err + } + return tx.TouchActivity() + }) +} + +// LifecycleNode routes a hosted Environment to the node lifecycle that +// provisions it, before its allocation exists: empty without a node. An +// existing allocation must agree with the Session's immutable placement. +func (s *Service) LifecycleNode(ctx context.Context, tenant, environment string) (string, error) { + key, err := AllocationKey{TenantID: tenant, EnvironmentID: environment}.parse() + if err != nil { + return "", err + } + p, err := s.reader.LifecyclePlacement(ctx, key) + if err != nil { + return "", err + } + if p.Provider == "" || p.Mode == "direct" { + if p.PlacementNodeID != "" || p.AllocationNodeID != "" { + return "", placement.ErrNodeUnavailable + } + return "", nil + } + if p.PlacementNodeID == "" || (p.AllocationID != "" && p.AllocationNodeID != p.PlacementNodeID) || (p.AllocationID == "" && p.PlacementReleased) { + return "", placement.ErrNodeUnavailable + } + return p.PlacementNodeID, nil +} + +// AllocationGeneration returns the node and generation an unreleased +// allocation was created on, including for a deleted Session. It never +// substitutes the deployment's target. +func (s *Service) AllocationGeneration(ctx context.Context, ref sandbox.Reference) (string, uint64, error) { + a, err := s.reader.EnvironmentAllocation(ctx, AllocationKey{TenantID: ref.TenantID, EnvironmentID: ref.EnvironmentID}) + if err != nil { + return "", 0, err + } + if a.State == "released" || a.ID != ref.AllocationID || a.NodeID == "" || a.DeploymentGeneration == 0 { + return "", 0, placement.ErrNodeUnavailable + } + return a.NodeID, a.DeploymentGeneration, nil +} diff --git a/services/core/internal/deployment/allocations_test.go b/services/core/internal/deployment/allocations_test.go new file mode 100644 index 000000000..cdb44204d --- /dev/null +++ b/services/core/internal/deployment/allocations_test.go @@ -0,0 +1,503 @@ +package deployment + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "slices" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +type fakeReservationTx struct { + t testing.TB + loadBoundDevice func() (bool, error) + insertEnvironmentDevice func(sessions.ExecutionDevice, string) error + loadEnvironment func() (sessions.Environment, error) + findAllocation func() (Allocation, bool, error) + lockDeployment func() (placement.Deployment, error) + loadReserved func() (placement.Reserved, error) + insertAllocation func(NewAllocation) (Allocation, error) +} + +func (f *fakeReservationTx) LoadBoundDevice(context.Context) (bool, error) { + if f.loadBoundDevice == nil { + unexpected(f.t, "LoadBoundDevice") + } + return f.loadBoundDevice() +} + +func (f *fakeReservationTx) InsertEnvironmentDevice(_ context.Context, device sessions.ExecutionDevice, credentialHash string) error { + if f.insertEnvironmentDevice == nil { + unexpected(f.t, "InsertEnvironmentDevice") + } + return f.insertEnvironmentDevice(device, credentialHash) +} + +func (f *fakeReservationTx) LoadEnvironment(context.Context) (sessions.Environment, error) { + if f.loadEnvironment == nil { + unexpected(f.t, "LoadEnvironment") + } + return f.loadEnvironment() +} + +func (f *fakeReservationTx) FindAllocation() (Allocation, bool, error) { + if f.findAllocation == nil { + unexpected(f.t, "FindAllocation") + } + return f.findAllocation() +} + +func (f *fakeReservationTx) LockDeployment() (placement.Deployment, error) { + if f.lockDeployment == nil { + unexpected(f.t, "LockDeployment") + } + return f.lockDeployment() +} + +func (f *fakeReservationTx) LoadReserved() (placement.Reserved, error) { + if f.loadReserved == nil { + unexpected(f.t, "LoadReserved") + } + return f.loadReserved() +} + +func (f *fakeReservationTx) InsertAllocation(allocation NewAllocation) (Allocation, error) { + if f.insertAllocation == nil { + unexpected(f.t, "InsertAllocation") + } + return f.insertAllocation(allocation) +} + +type fakeAllocationTx struct { + t testing.TB + loadAllocation func() (Allocation, error) + loadSessionDevice func() (SessionDevice, bool, error) + settleCreation func(Allocation) (Allocation, error) + keep func(Allocation) (Allocation, error) + release func(Allocation) (Allocation, error) +} + +func (f *fakeAllocationTx) LoadAllocation() (Allocation, error) { + if f.loadAllocation == nil { + unexpected(f.t, "LoadAllocation") + } + return f.loadAllocation() +} + +func (f *fakeAllocationTx) LoadSessionDevice() (SessionDevice, bool, error) { + if f.loadSessionDevice == nil { + unexpected(f.t, "LoadSessionDevice") + } + return f.loadSessionDevice() +} + +func (f *fakeAllocationTx) LoadActivity(Allocation) (Activity, error) { + unexpected(f.t, "LoadActivity") + return Activity{}, nil +} + +func (f *fakeAllocationTx) LoadRestore(Allocation) (placement.Restore, error) { + unexpected(f.t, "LoadRestore") + return placement.Restore{}, nil +} + +func (f *fakeAllocationTx) ObserveRunning(Allocation) (Allocation, error) { + unexpected(f.t, "ObserveRunning") + return Allocation{}, nil +} + +func (f *fakeAllocationTx) Keep(current Allocation) (Allocation, error) { + if f.keep == nil { + unexpected(f.t, "Keep") + } + return f.keep(current) +} + +func (f *fakeAllocationTx) SettleCreation(current Allocation) (Allocation, error) { + if f.settleCreation == nil { + unexpected(f.t, "SettleCreation") + } + return f.settleCreation(current) +} + +func (f *fakeAllocationTx) Release(current Allocation) (Allocation, error) { + if f.release == nil { + unexpected(f.t, "Release") + } + return f.release(current) +} + +func (f *fakeAllocationTx) SetCompute(Allocation, ComputeChange) (Allocation, error) { + unexpected(f.t, "SetCompute") + return Allocation{}, nil +} + +func (f *fakeAllocationTx) RecordObservation(Allocation, string) error { + unexpected(f.t, "RecordObservation") + return nil +} + +// fakeCleanupTx is a strict allocation cleanup transaction. Its Session +// methods cover a Session without an active Turn, pending input or input +// activity change. +type fakeCleanupTx struct { + *fakeAllocationTx + revokeDevice func(Allocation) error + requestCleanup func(Allocation) (Allocation, error) + loadEnvironment func() (sessions.Environment, error) + loadEnvironmentInput func() (*sessions.EnvironmentInputState, error) + expireEnvironment func(string) error + failPendingInput func() error + loadActiveTurn func() (sessions.Turn, bool, error) + cancelPendingInput func() error +} + +func (f *fakeCleanupTx) RevokeDevice(current Allocation) error { + if f.revokeDevice == nil { + unexpected(f.t, "RevokeDevice") + } + return f.revokeDevice(current) +} + +func (f *fakeCleanupTx) RequestCleanup(current Allocation) (Allocation, error) { + if f.requestCleanup == nil { + unexpected(f.t, "RequestCleanup") + } + return f.requestCleanup(current) +} + +func (f *fakeCleanupTx) LoadEnvironment(context.Context) (sessions.Environment, error) { + if f.loadEnvironment == nil { + unexpected(f.t, "LoadEnvironment") + } + return f.loadEnvironment() +} + +func (f *fakeCleanupTx) LoadEnvironmentInput(context.Context) (*sessions.EnvironmentInputState, error) { + if f.loadEnvironmentInput == nil { + unexpected(f.t, "LoadEnvironmentInput") + } + return f.loadEnvironmentInput() +} + +func (f *fakeCleanupTx) ExpireEnvironment(_ context.Context, environment string) error { + if f.expireEnvironment == nil { + unexpected(f.t, "ExpireEnvironment") + } + return f.expireEnvironment(environment) +} + +func (f *fakeCleanupTx) FailPendingInput(context.Context) error { + if f.failPendingInput == nil { + unexpected(f.t, "FailPendingInput") + } + return f.failPendingInput() +} + +func (f *fakeCleanupTx) LoadActiveTurn(context.Context) (sessions.Turn, bool, error) { + if f.loadActiveTurn == nil { + unexpected(f.t, "LoadActiveTurn") + } + return f.loadActiveTurn() +} + +func (f *fakeCleanupTx) CancelPendingInput(context.Context) error { + if f.cancelPendingInput == nil { + unexpected(f.t, "CancelPendingInput") + } + return f.cancelPendingInput() +} + +func (f *fakeCleanupTx) RecordEnvironmentFailure(context.Context, string, string, *sessions.ProvisioningFailureDetail) (time.Time, error) { + unexpected(f.t, "RecordEnvironmentFailure") + return time.Time{}, nil +} + +func (f *fakeCleanupTx) RequestTurnCancel(context.Context, string) error { + unexpected(f.t, "RequestTurnCancel") + return nil +} + +func (f *fakeCleanupTx) LoadTurn(context.Context, string) (sessions.Turn, error) { + unexpected(f.t, "LoadTurn") + return sessions.Turn{}, nil +} + +func (f *fakeCleanupTx) LoadEnding(context.Context, string) (sessions.Ending, error) { + unexpected(f.t, "LoadEnding") + return sessions.Ending{}, nil +} + +func (f *fakeCleanupTx) ApplyTurnEnd(context.Context, string, sessions.TurnEnd) error { + unexpected(f.t, "ApplyTurnEnd") + return nil +} + +func (f *fakeCleanupTx) LoadUsage(context.Context) (json.RawMessage, error) { + unexpected(f.t, "LoadUsage") + return nil, nil +} + +func (f *fakeCleanupTx) AppendChanges(context.Context, ...sessions.SessionChange) error { + unexpected(f.t, "AppendChanges") + return nil +} + +// allocationOperations builds execution operations whose reservations and +// allocation changes run on the given transactions. +func allocationOperations(t *testing.T, reservation *fakeReservationTx, locked sessions.LockedSession, allocation *fakeAllocationTx) *ExecutionOperations { + t.Helper() + storage := &fakeExecutionStorage{t: t} + if reservation != nil { + storage.withReservation = func(_ context.Context, _ AllocationKey, apply func(sessions.LockedSession, ReservationTx) error) error { + return apply(locked, reservation) + } + } + if allocation != nil { + storage.withAllocation = func(_ context.Context, _ AllocationKey, apply func(AllocationTx) error) error { + return apply(allocation) + } + } + result, err := NewExecutionOperations(newService(t, &fakeStorage{t: t}, &fakeReader{t: t}, testPublicURL), storage) + if err != nil { + t.Fatal(err) + } + return result +} + +func testCredentialHash() string { + digest := sha256.Sum256([]byte("credential")) + return hex.EncodeToString(digest[:]) +} + +func hostedEnvironment(id string) func() (sessions.Environment, error) { + return func() (sessions.Environment, error) { + return sessions.Environment{ID: id, Status: "pending", Configuration: json.RawMessage(`{"type":"openai_hosted"}`)}, nil + } +} + +// An existing allocation replays before admission is read, so retries and +// cleanup continue while admission is closed; another installation conflicts. +func TestReserveAllocationReplaysBeforeAdmission(t *testing.T) { + key := AllocationKey{TenantID: uuid.NewString(), EnvironmentID: uuid.NewString()} + installation := uuid.NewString() + existing := Allocation{ID: uuid.NewString(), EnvironmentID: key.EnvironmentID, ProviderKey: installation} + tx := &fakeReservationTx{t: t, loadEnvironment: hostedEnvironment(key.EnvironmentID), + findAllocation: func() (Allocation, bool, error) { return existing, true, nil }} + replayed, err := allocationOperations(t, tx, sessions.LockedSession{}, nil).ReserveAllocation(t.Context(), key, installation, testCredentialHash()) + if err != nil || !replayed.Replayed || replayed.ID != existing.ID { + t.Fatal("reservation did not replay", replayed, err) + } + if _, err := allocationOperations(t, tx, sessions.LockedSession{}, nil).ReserveAllocation(t.Context(), key, uuid.NewString(), testCredentialHash()); !errors.Is(err, ErrAllocationConflict) { + t.Fatal("another installation replayed the allocation", err) + } + deleted := &fakeReservationTx{t: t} + if _, err := allocationOperations(t, deleted, sessions.LockedSession{Deleted: true}, nil).ReserveAllocation(t.Context(), key, installation, testCredentialHash()); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal("a deleted Session reserved an allocation", err) + } + selfHosted := &fakeReservationTx{t: t, loadEnvironment: func() (sessions.Environment, error) { + return sessions.Environment{ID: key.EnvironmentID, Configuration: json.RawMessage(`{"type":"self_hosted"}`)}, nil + }} + if _, err := allocationOperations(t, selfHosted, sessions.LockedSession{}, nil).ReserveAllocation(t.Context(), key, installation, testCredentialHash()); !errors.Is(err, ErrInvalidInput) { + t.Fatal("a self-hosted Environment reserved an allocation", err) + } +} + +// A fresh reservation passes admission, takes the node and generation its +// Session reserved and creates the dedicated device with the allocation. +func TestReserveAllocationAdmitsAndTakesTheReservedNode(t *testing.T) { + key := AllocationKey{TenantID: uuid.NewString(), EnvironmentID: uuid.NewString()} + installation, node := uuid.NewString(), uuid.NewString() + deployment := placement.Deployment{InstallationID: installation, Mode: "nodes", Generation: 7} + fresh := func() *fakeReservationTx { + return &fakeReservationTx{t: t, loadEnvironment: hostedEnvironment(key.EnvironmentID), + findAllocation: func() (Allocation, bool, error) { return Allocation{}, false, nil }, + lockDeployment: func() (placement.Deployment, error) { return deployment, nil }} + } + paused := fresh() + paused.lockDeployment = func() (placement.Deployment, error) { + d := deployment + d.AdmissionPaused = true + return d, nil + } + if _, err := allocationOperations(t, paused, sessions.LockedSession{}, nil).ReserveAllocation(t.Context(), key, installation, testCredentialHash()); !errors.Is(err, placement.ErrAdmissionClosed) { + t.Fatal("paused admission reserved an allocation", err) + } + released := fresh() + released.loadReserved = func() (placement.Reserved, error) { + return placement.Reserved{NodeID: node, Generation: 5, Released: true, Available: true}, nil + } + if _, err := allocationOperations(t, released, sessions.LockedSession{}, nil).ReserveAllocation(t.Context(), key, installation, testCredentialHash()); !errors.Is(err, placement.ErrNodeUnavailable) { + t.Fatal("a released placement reserved an allocation", err) + } + var device sessions.ExecutionDevice + var inserted NewAllocation + tx := fresh() + tx.loadReserved = func() (placement.Reserved, error) { + return placement.Reserved{NodeID: node, Generation: 5, Available: true}, nil + } + tx.loadBoundDevice = func() (bool, error) { return false, nil } + tx.insertEnvironmentDevice = func(d sessions.ExecutionDevice, hash string) error { + if hash != testCredentialHash() { + t.Fatal("device credential", hash) + } + device = d + return nil + } + tx.insertAllocation = func(a NewAllocation) (Allocation, error) { + inserted = a + return Allocation{ID: a.ID, DeviceID: a.DeviceID, NodeID: a.NodeID}, nil + } + result, err := allocationOperations(t, tx, sessions.LockedSession{}, nil).ReserveAllocation(t.Context(), key, installation, testCredentialHash()) + if err != nil || result.Replayed || inserted.NodeID != node || inserted.Generation != 5 || inserted.ProviderKey != installation || inserted.DeviceID != device.ID || device.EnvironmentID != key.EnvironmentID { + t.Fatal("reservation", result, inserted, device, err) + } +} + +// A live change needs the Session undeleted and bound to the allocation's +// device; settlement continues for a deleted Session. Any other owner is a +// conflict. +func TestAllocationChangesCheckTheOwner(t *testing.T) { + owner := Allocation{ID: uuid.NewString(), DeviceID: uuid.NewString(), EnvironmentID: uuid.NewString(), TenantID: uuid.NewString(), ProviderKey: uuid.NewString()} + stored := func(current Allocation) func() (Allocation, error) { + return func() (Allocation, error) { return current, nil } + } + moved := owner + moved.DeviceID = uuid.NewString() + if _, err := allocationOperations(t, nil, sessions.LockedSession{}, &fakeAllocationTx{t: t, loadAllocation: stored(moved)}).SettleCreation(t.Context(), owner); !errors.Is(err, ErrAllocationConflict) { + t.Fatal("a replaced allocation settled", err) + } + deleted := owner + deleted.SessionDeleted = true + if _, err := allocationOperations(t, nil, sessions.LockedSession{}, &fakeAllocationTx{t: t, loadAllocation: stored(deleted)}).KeepAllocation(t.Context(), owner); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal("a deleted Session kept its allocation", err) + } + settled := &fakeAllocationTx{t: t, loadAllocation: stored(deleted), settleCreation: func(current Allocation) (Allocation, error) { + current.CreateSettled = true + return current, nil + }} + if result, err := allocationOperations(t, nil, sessions.LockedSession{}, settled).SettleCreation(t.Context(), owner); err != nil || !result.CreateSettled { + t.Fatal("a deleted Session's creation did not settle", result, err) + } + unbound := &fakeAllocationTx{t: t, loadAllocation: stored(owner), loadSessionDevice: func() (SessionDevice, bool, error) { + return SessionDevice{ID: uuid.NewString(), EnvironmentID: owner.EnvironmentID}, true, nil + }} + if _, err := allocationOperations(t, nil, sessions.LockedSession{}, unbound).KeepAllocation(t.Context(), owner); !errors.Is(err, ErrAllocationConflict) { + t.Fatal("a Session bound to another device kept the allocation", err) + } +} + +// SetCompute rejects an invalid phase change before it reaches storage. +func TestSetComputeValidatesBeforeStorage(t *testing.T) { + owner := Allocation{ID: uuid.NewString(), TenantID: uuid.NewString(), EnvironmentID: uuid.NewString(), ComputePhase: "running"} + until := time.Now().Add(time.Hour) + for name, call := range map[string]func(*ExecutionOperations) error{ + "skipped phase": func(o *ExecutionOperations) error { + _, err := o.SetCompute(t.Context(), owner, "suspended", json.RawMessage(`{}`), &until, time.Second) + return err + }, + "state not object": func(o *ExecutionOperations) error { + _, err := o.SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`[]`), &until, time.Second) + return err + }, + "no retention": func(o *ExecutionOperations) error { + _, err := o.SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), nil, time.Second) + return err + }, + "no idle timeout": func(o *ExecutionOperations) error { + _, err := o.SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, 0) + return err + }, + } { + if err := call(allocationOperations(t, nil, sessions.LockedSession{}, nil)); !errors.Is(err, ErrInvalidInput) { + t.Fatal(name, err) + } + } +} + +// A malformed credential digest is deployment's invalid input and never +// reaches storage. +func TestReserveAllocationValidatesTheCredentialDigest(t *testing.T) { + key := AllocationKey{TenantID: uuid.NewString(), EnvironmentID: uuid.NewString()} + if _, err := allocationOperations(t, nil, sessions.LockedSession{}, nil).ReserveAllocation(t.Context(), key, uuid.NewString(), "not-a-digest"); !errors.Is(err, ErrInvalidInput) { + t.Fatal("a malformed digest reserved an allocation", err) + } +} + +// Cleanup revokes the device, then cancels a deleted Session's work or +// terminates a live Session's Environment, then requests cleanup. An absent +// creation then settles and releases the allocation, which releases its node +// placement with it. +func TestCleanupRevokesSettlesTheSessionThenReleases(t *testing.T) { + owner := Allocation{ID: uuid.NewString(), DeviceID: uuid.NewString(), EnvironmentID: uuid.NewString(), TenantID: uuid.NewString(), ProviderKey: uuid.NewString(), State: "running"} + for _, test := range []struct { + name string + deleted, absent bool + want []string + }{ + {"deleted Session with absent creation", true, true, []string{"LoadAllocation", "RevokeDevice", "LoadAllocation", "LoadActiveTurn", "CancelPendingInput", "RequestCleanup", "SettleCreation", "Release"}}, + {"live Session with expired compute", false, false, []string{"LoadAllocation", "RevokeDevice", "LoadAllocation", "LoadEnvironment", "LoadEnvironmentInput", "ExpireEnvironment", "FailPendingInput", "LoadActiveTurn", "CancelPendingInput", "LoadEnvironmentInput", "RequestCleanup"}}, + } { + t.Run(test.name, func(t *testing.T) { + var calls []string + stored := owner + stored.SessionDeleted, stored.Expired = test.deleted, !test.deleted + change := func(name string, apply func(*Allocation)) func(Allocation) (Allocation, error) { + return func(current Allocation) (Allocation, error) { + calls = append(calls, name) + apply(¤t) + return current, nil + } + } + tx := &fakeCleanupTx{ + fakeAllocationTx: &fakeAllocationTx{t: t, + loadAllocation: func() (Allocation, error) { calls = append(calls, "LoadAllocation"); return stored, nil }, + settleCreation: change("SettleCreation", func(a *Allocation) { a.CreateSettled = true }), + release: change("Release", func(a *Allocation) { a.State = "released" }), + }, + revokeDevice: func(Allocation) error { calls = append(calls, "RevokeDevice"); return nil }, + requestCleanup: change("RequestCleanup", func(a *Allocation) { a.State = "cleanup_pending" }), + loadActiveTurn: func() (sessions.Turn, bool, error) { + calls = append(calls, "LoadActiveTurn") + return sessions.Turn{}, false, nil + }, + cancelPendingInput: func() error { calls = append(calls, "CancelPendingInput"); return nil }, + } + if !test.deleted { + tx.loadEnvironment = func() (sessions.Environment, error) { + calls = append(calls, "LoadEnvironment") + return sessions.Environment{ID: owner.EnvironmentID, Status: "ready"}, nil + } + tx.loadEnvironmentInput = func() (*sessions.EnvironmentInputState, error) { + calls = append(calls, "LoadEnvironmentInput") + return nil, nil + } + tx.expireEnvironment = func(string) error { calls = append(calls, "ExpireEnvironment"); return nil } + tx.failPendingInput = func() error { calls = append(calls, "FailPendingInput"); return nil } + } + storage := &fakeExecutionStorage{t: t, withAllocationCleanup: func(_ context.Context, _ AllocationKey, apply func(AllocationCleanupTx) error) error { + return apply(tx) + }} + operations, err := NewExecutionOperations(newService(t, &fakeStorage{t: t}, &fakeReader{t: t}, testPublicURL), storage) + if err != nil { + t.Fatal(err) + } + cleanup, want := operations.RequestCleanup, "cleanup_pending" + if test.absent { + cleanup, want = operations.ReleaseAbsentCreation, "released" + } + result, err := cleanup(t.Context(), owner) + if err != nil || result.State != want || !slices.Equal(calls, test.want) { + t.Fatal("cleanup order", result.State, calls, err) + } + }) + } +} diff --git a/services/core/internal/deployment/errors.go b/services/core/internal/deployment/errors.go index 3d6c6a1f8..27252d080 100644 --- a/services/core/internal/deployment/errors.go +++ b/services/core/internal/deployment/errors.go @@ -18,14 +18,13 @@ var ( ErrSpecificationMismatch = errors.New("node does not match the deployment specification") ErrResetInProgress = errors.New("a sandbox reset is in progress") ErrNotConfigured = errors.New("the sandbox deployment is not configured") - // ErrPublicURLUnreachable rejects selection and admission when the provider - // requires a reachable public origin and the installation is loopback. - ErrPublicURLUnreachable = errors.New("This sandbox provider needs a reachable HTTPS public URL before they can connect to Core.") - ErrNodesPreparing = errors.New("sandbox nodes are preparing the target generation") - ErrNodeUnavailable = errors.New("sandbox node unavailable") - ErrNodeInUse = errors.New("sandbox node retains resources") - ErrNodeCredential = errors.New("invalid sandbox node credential") - ErrLocalNodeConfigured = errors.New("local sandbox node is enabled in deployment configuration") + ErrNodeInUse = errors.New("sandbox node retains resources") + ErrNodeCredential = errors.New("invalid sandbox node credential") + ErrLocalNodeConfigured = errors.New("local sandbox node is enabled in deployment configuration") + // ErrAllocationConflict rejects an allocation change whose owner no longer + // matches the stored allocation, device binding, state or compute revision, + // or a replay for another installation. + ErrAllocationConflict = errors.New("the sandbox allocation changed") // ErrNodeAddressMismatch rejects an enrollment whose Core address is not // the installation public URL. The token stays unconsumed. ErrNodeAddressMismatch = errors.New("sandbox node Core address differs from the public URL") diff --git a/services/core/internal/deployment/fakes_test.go b/services/core/internal/deployment/fakes_test.go index bdd5ee747..8113024ec 100644 --- a/services/core/internal/deployment/fakes_test.go +++ b/services/core/internal/deployment/fakes_test.go @@ -9,6 +9,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/adminaudit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/coremetrics" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) // The fakes are strict: each method runs its func field, and a nil field fails @@ -25,6 +26,7 @@ type fakeStorage struct { connectNode func(context.Context, string, string, uint64) (bool, error) disconnectNode func(context.Context, string, string, uint64) error sampleHostHistory func(context.Context) (int64, error) + withActivity func(context.Context, AllocationKey, func(sessions.LockedSession, ActivityTx) error) error } func (f *fakeStorage) WithNodes(ctx context.Context, apply func(NodeTx) error) error { @@ -56,8 +58,12 @@ func (f *fakeStorage) SampleHostHistory(ctx context.Context) (int64, error) { } type fakeExecutionStorage struct { - t testing.TB - withDeployment func(context.Context, func(DeploymentTx) error) error + t testing.TB + withDeployment func(context.Context, func(DeploymentTx) error) error + withReservation func(context.Context, AllocationKey, func(sessions.LockedSession, ReservationTx) error) error + withAllocation func(context.Context, AllocationKey, func(AllocationTx) error) error + withAllocationCleanup func(context.Context, AllocationKey, func(AllocationCleanupTx) error) error + clearWake func(context.Context, string, time.Time) error } func (f *fakeExecutionStorage) WithDeployment(ctx context.Context, apply func(DeploymentTx) error) error { @@ -68,17 +74,29 @@ func (f *fakeExecutionStorage) WithDeployment(ctx context.Context, apply func(De } type fakeReader struct { - t testing.TB - deployment func(context.Context) (Record, error) - snapshot func(context.Context) (Snapshot, error) - ownerEpoch func(context.Context) (uint64, error) - allocation func(context.Context, sandbox.Reference) (AllocationRecord, error) - generations func(context.Context, int64) ([]GenerationRecord, error) - nodes func(context.Context) ([]NodeRecord, error) - nodeHistory func(context.Context, string, coremetrics.Range) (NodeRecord, []HostHistoryPoint, error) - readNodes func(context.Context, func(NodeReads) error) error - resetSessions func(context.Context, string, bool) ([]ResetSession, error) - addressBindings func(context.Context, string) (AddressBindings, error) + t testing.TB + deployment func(context.Context) (Record, error) + snapshot func(context.Context) (Snapshot, error) + ownerEpoch func(context.Context) (uint64, error) + allocation func(context.Context, sandbox.Reference) (AllocationRecord, error) + generations func(context.Context, int64) ([]GenerationRecord, error) + nodes func(context.Context) ([]NodeRecord, error) + nodeHistory func(context.Context, string, coremetrics.Range) (NodeRecord, []HostHistoryPoint, error) + readNodes func(context.Context, func(NodeReads) error) error + resetSessions func(context.Context, string, bool) ([]ResetSession, error) + addressBindings func(context.Context, string) (AddressBindings, error) + environmentAllocation func(context.Context, AllocationKey) (Allocation, error) + credentialAllocations func(context.Context, string) ([]Allocation, error) + observationSessions func(context.Context, string, int) (ObservationSessionPage, error) + nodeAllocations func(context.Context, string) ([]NodeAllocation, error) + nodeOnline func(context.Context, string) (bool, error) + lifecycleNodes func(context.Context) ([]string, error) + lifecycleAllocations func(context.Context, string, string) ([]Allocation, error) + unallocatedEnvironments func(context.Context, string, string) ([]UnallocatedEnvironment, error) + lifecyclePlacement func(context.Context, AllocationKey) (LifecyclePlacement, error) + activity func(context.Context, string) (Activity, error) + countComputeReservations func(context.Context, string) (int64, error) + countRetainedAllocations func(context.Context, string) (int64, error) } func (f *fakeReader) Deployment(ctx context.Context) (Record, error) { @@ -490,3 +508,122 @@ func (f *fakeDeploymentTx) RecordAuditAs(source adminaudit.Source, action, insta } return f.recordAuditAs(source, action, installationID) } + +func (f *fakeStorage) WithActivity(ctx context.Context, key AllocationKey, apply func(sessions.LockedSession, ActivityTx) error) error { + if f.withActivity == nil { + unexpected(f.t, "WithActivity") + } + return f.withActivity(ctx, key, apply) +} + +func (f *fakeExecutionStorage) WithReservation(ctx context.Context, key AllocationKey, apply func(sessions.LockedSession, ReservationTx) error) error { + if f.withReservation == nil { + unexpected(f.t, "WithReservation") + } + return f.withReservation(ctx, key, apply) +} + +func (f *fakeExecutionStorage) WithAllocation(ctx context.Context, key AllocationKey, apply func(AllocationTx) error) error { + if f.withAllocation == nil { + unexpected(f.t, "WithAllocation") + } + return f.withAllocation(ctx, key, apply) +} + +func (f *fakeExecutionStorage) WithAllocationCleanup(ctx context.Context, key AllocationKey, apply func(AllocationCleanupTx) error) error { + if f.withAllocationCleanup == nil { + unexpected(f.t, "WithAllocationCleanup") + } + return f.withAllocationCleanup(ctx, key, apply) +} + +func (f *fakeExecutionStorage) ClearWake(ctx context.Context, allocationID string, observed time.Time) error { + if f.clearWake == nil { + unexpected(f.t, "ClearWake") + } + return f.clearWake(ctx, allocationID, observed) +} + +func (f *fakeReader) EnvironmentAllocation(ctx context.Context, key AllocationKey) (Allocation, error) { + if f.environmentAllocation == nil { + unexpected(f.t, "EnvironmentAllocation") + } + return f.environmentAllocation(ctx, key) +} + +func (f *fakeReader) CredentialAllocations(ctx context.Context, after string) ([]Allocation, error) { + if f.credentialAllocations == nil { + unexpected(f.t, "CredentialAllocations") + } + return f.credentialAllocations(ctx, after) +} + +func (f *fakeReader) ObservationSessions(ctx context.Context, after string, limit int) (ObservationSessionPage, error) { + if f.observationSessions == nil { + unexpected(f.t, "ObservationSessions") + } + return f.observationSessions(ctx, after, limit) +} + +func (f *fakeReader) NodeAllocations(ctx context.Context, nodeID string) ([]NodeAllocation, error) { + if f.nodeAllocations == nil { + unexpected(f.t, "NodeAllocations") + } + return f.nodeAllocations(ctx, nodeID) +} + +func (f *fakeReader) NodeOnline(ctx context.Context, nodeID string) (bool, error) { + if f.nodeOnline == nil { + unexpected(f.t, "NodeOnline") + } + return f.nodeOnline(ctx, nodeID) +} + +func (f *fakeReader) LifecycleNodes(ctx context.Context) ([]string, error) { + if f.lifecycleNodes == nil { + unexpected(f.t, "LifecycleNodes") + } + return f.lifecycleNodes(ctx) +} + +func (f *fakeReader) LifecycleAllocations(ctx context.Context, nodeID string, after string) ([]Allocation, error) { + if f.lifecycleAllocations == nil { + unexpected(f.t, "LifecycleAllocations") + } + return f.lifecycleAllocations(ctx, nodeID, after) +} + +func (f *fakeReader) UnallocatedEnvironments(ctx context.Context, nodeID string, after string) ([]UnallocatedEnvironment, error) { + if f.unallocatedEnvironments == nil { + unexpected(f.t, "UnallocatedEnvironments") + } + return f.unallocatedEnvironments(ctx, nodeID, after) +} + +func (f *fakeReader) LifecyclePlacement(ctx context.Context, key AllocationKey) (LifecyclePlacement, error) { + if f.lifecyclePlacement == nil { + unexpected(f.t, "LifecyclePlacement") + } + return f.lifecyclePlacement(ctx, key) +} + +func (f *fakeReader) Activity(ctx context.Context, allocationID string) (Activity, error) { + if f.activity == nil { + unexpected(f.t, "Activity") + } + return f.activity(ctx, allocationID) +} + +func (f *fakeReader) CountComputeReservations(ctx context.Context, installationID string) (int64, error) { + if f.countComputeReservations == nil { + unexpected(f.t, "CountComputeReservations") + } + return f.countComputeReservations(ctx, installationID) +} + +func (f *fakeReader) CountRetainedAllocations(ctx context.Context, installationID string) (int64, error) { + if f.countRetainedAllocations == nil { + unexpected(f.t, "CountRetainedAllocations") + } + return f.countRetainedAllocations(ctx, installationID) +} diff --git a/services/core/internal/deployment/nodes.go b/services/core/internal/deployment/nodes.go index 4feef8603..3f2ab3764 100644 --- a/services/core/internal/deployment/nodes.go +++ b/services/core/internal/deployment/nodes.go @@ -10,6 +10,7 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/coremetrics" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/google/uuid" ) @@ -79,7 +80,7 @@ func (s *Service) withManager(ctx context.Context, apply func(NodeTx, Record) er return err } if !initialized(d) { - return ErrNodeUnavailable + return placement.ErrNodeUnavailable } return apply(tx, d) }) @@ -199,7 +200,7 @@ func (s *Service) Enroll(ctx context.Context, token string, input Enrollment) (N return ErrNodeCredential } if err != nil { - return ErrNodeUnavailable + return placement.ErrNodeUnavailable } if receipt.Consumed || !receipt.ExpiresAt.After(time.Now()) { return ErrNodeCredential @@ -208,7 +209,7 @@ func (s *Service) Enroll(ctx context.Context, token string, input Enrollment) (N return ErrNodeCredential } if !initialized(d) { - return ErrNodeUnavailable + return placement.ErrNodeUnavailable } if d.Reset != nil { return ErrResetInProgress @@ -223,7 +224,7 @@ func (s *Service) Enroll(ctx context.Context, token string, input Enrollment) (N // The node must use the address Core advertises now. It read that address // from its configuration, but the public URL may have changed since, or an // operator may have registered by hand with another origin. - if input.CoreURL != s.publicURL { + if input.CoreURL != s.rules.PublicURL() { return ErrNodeAddressMismatch } retained, err := s.registry.RetainedLimit(d.Provider, receipt.MaxActive, receipt.MaxRetained) @@ -271,14 +272,14 @@ func (s *Service) AuthenticateNode(ctx context.Context, nodeID, credential strin return ErrNodeCredential } if err != nil { - return ErrNodeUnavailable + return placement.ErrNodeUnavailable } if n.CredentialDigest != tokenDigest(credential) { return ErrNodeCredential } d, err := tx.LoadDeployment() if err != nil { - return ErrNodeUnavailable + return placement.ErrNodeUnavailable } if n.InstallationID != d.InstallationID || d.Mode != "nodes" { return ErrNodeCredential @@ -394,7 +395,7 @@ func (s *Service) NodeConfiguration(ctx context.Context, nodeID, token string, g return ErrNodeCredential } if !initialized(d) { - return ErrNodeUnavailable + return placement.ErrNodeUnavailable } if node == nil && d.Reset != nil { return ErrResetInProgress @@ -432,7 +433,7 @@ func (s *Service) NodeConfiguration(ctx context.Context, nodeID, token string, g if err != nil { return err } - result = NodeConfiguration{MaxActive: active, MaxRetained: limit, InstallationID: d.InstallationID, Provider: d.Provider, CoreURL: s.publicURL, Generation: selected, Specification: spec, SpecificationDigest: spec.Digest(d.Provider)} + result = NodeConfiguration{MaxActive: active, MaxRetained: limit, InstallationID: d.InstallationID, Provider: d.Provider, CoreURL: s.rules.PublicURL(), Generation: selected, Specification: spec, SpecificationDigest: spec.Digest(d.Provider)} return nil }) if err != nil { diff --git a/services/core/internal/deployment/placement/placement.go b/services/core/internal/deployment/placement/placement.go new file mode 100644 index 000000000..433e6f6b8 --- /dev/null +++ b/services/core/internal/deployment/placement/placement.go @@ -0,0 +1,244 @@ +// Package placement decides whether the sandbox deployment admits new hosted +// work and which node receives it. Its rules are pure: they read the +// deployment and node facts the caller loads under the deployment lock, the +// provider declarations and the installation public URL, and do no I/O. The +// caller applies what they decide in the same transaction. +package placement + +import ( + "encoding/json" + "errors" + "fmt" + "net" + "net/url" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" +) + +var ( + // ErrResetAdmission rejects new hosted work while a sandbox reset pauses + // admission. + ErrResetAdmission = errors.New("hosted admission is paused for a sandbox reset") + // ErrAdmissionClosed rejects new hosted work that the deployment cannot + // admit: paused for maintenance, without a valid specification, or for + // another installation. + ErrAdmissionClosed = errors.New("environment is no longer available") + // ErrPublicURLUnreachable rejects selection and admission when the provider + // requires a reachable public origin and the installation is loopback. + ErrPublicURLUnreachable = errors.New("This sandbox provider needs a reachable HTTPS public URL before they can connect to Core.") + // ErrNodesPreparing rejects placement while no node serves the target + // generation and at least one is preparing it. + ErrNodesPreparing = errors.New("sandbox nodes are preparing the target generation") + // ErrNodeUnavailable reports a sandbox node that is offline, unready or + // full, or no node at all. + ErrNodeUnavailable = errors.New("sandbox node unavailable") +) + +// Declarations are the provider declarations placement reads. +// *providers.Registry satisfies it. +type Declarations interface { + // RequiresPublicOrigin reports whether the provider's guests must reach + // Core at a public origin. + RequiresPublicOrigin(provider string) (bool, error) + // ValidateSpecification rejects a specification the provider cannot run. + ValidateSpecification(provider string, spec sandbox.DeploymentSpec) error +} + +// Rules decides admission and placement for one installation. Its +// declarations and public URL never change after NewRules. +type Rules struct { + declarations Declarations + publicURL string +} + +// NewRules returns the rules for the provider declarations and the +// installation public URL, which is empty when the installation has none. +func NewRules(declarations Declarations, publicURL string) (*Rules, error) { + if declarations == nil { + return nil, errors.New("placement rules require provider declarations") + } + return &Rules{declarations: declarations, publicURL: publicURL}, nil +} + +// PublicURL returns the installation public URL, OAC_PUBLIC_URL. Core +// derives every address it gives nodes, sandboxes and administrators from it. +func (r *Rules) PublicURL() string { return r.publicURL } + +// Deployment is the deployment as placement reads it, loaded under the +// deployment lock. +type Deployment struct { + // InstallationID is empty until an installation is claimed or configured. + InstallationID string + Provider, Mode string + Generation uint64 + WebManaged bool + AdmissionPaused bool + // Resetting reports a sandbox reset in progress. + Resetting bool + Specification json.RawMessage +} + +// Node is one node of the installation with its presence and usage. +type Node struct { + ID string + Online, ServingReady bool + // ReadyGeneration is nil until the node serves a generation. + ReadyGeneration *uint64 + TargetState string + Active, Retained int64 + MaxActive, MaxRetained int + CoreURL string +} + +// Placement is the node and generation a new Session's Environment reserves. +type Placement struct { + NodeID string + Generation uint64 +} + +// Reserved is the node placement a Session reserved for its Environment. +type Reserved struct { + NodeID string + Generation uint64 + Released bool + // Available reports that the node is connected, fresh and ready for the + // reserved generation. + Available bool +} + +// Restore is what restoring suspended compute on its node reads: the node, +// nil when the installation no longer has it, and whether the node is ready +// for the allocation's generation. +type Restore struct { + Node *Node + // Generation is zero when the allocation has none. + Generation uint64 + GenerationReady bool +} + +// CheckPublicOrigin rejects a provider that requires a reachable public +// origin while the installation public URL is loopback. +func (r *Rules) CheckPublicOrigin(provider string) error { + required, err := r.declarations.RequiresPublicOrigin(provider) + if err != nil { + return err + } + if required && LoopbackOrigin(r.publicURL) { + return ErrPublicURLUnreachable + } + return nil +} + +// CheckAdmission admits new hosted work on the deployment. installation is +// the canonical installation the work was provisioned for, or empty for a new +// Session. An unclaimed deployment admits everything. +func (r *Rules) CheckAdmission(d Deployment, installation string) error { + if d.InstallationID == "" { + return nil + } + if d.Resetting { + return ErrResetAdmission + } + if d.AdmissionPaused { + return fmt.Errorf("%w: sandbox creation is paused for provider maintenance", ErrAdmissionClosed) + } + if d.Provider != "" { + var spec sandbox.DeploymentSpec + if json.Unmarshal(d.Specification, &spec) != nil || r.declarations.ValidateSpecification(d.Provider, spec) != nil { + return fmt.Errorf("%w: sandbox creation requires a deployment specification", ErrAdmissionClosed) + } + } + if installation != "" && installation != d.InstallationID { + return fmt.Errorf("%w: sandbox installation does not match deployment", ErrAdmissionClosed) + } + return nil +} + +// DecidePlacement chooses the node a new Session's hosted Environment +// reserves, or nil when the deployment places no node: in direct mode and +// without a provider. It prefers the highest ready generation, then the +// fewest active sandboxes. A Web-managed installation places only on nodes +// enrolled with its public URL; restores still reach the others. +func (r *Rules) DecidePlacement(d Deployment, nodes []Node) (*Placement, error) { + if d.Resetting { + return nil, ErrResetAdmission + } + // A changed installation address cannot admit guests that require a + // public origin. Existing owned resources remain available for cleanup. + if d.Provider != "" { + if err := r.CheckPublicOrigin(d.Provider); err != nil { + return nil, err + } + } + if d.Mode == "direct" { + if d.AdmissionPaused { + return nil, ErrNodeUnavailable + } + return nil, nil + } + if d.Provider == "" { + if d.WebManaged { + return nil, ErrNodeUnavailable + } + return nil, nil + } + if d.AdmissionPaused { + return nil, ErrNodeUnavailable + } + var chosen *Node + preparing := false + for i := range nodes { + n := &nodes[i] + free := n.Active < int64(n.MaxActive) && n.Retained < int64(n.MaxRetained) + reachable := !d.WebManaged || n.CoreURL == r.publicURL + if n.Online && n.TargetState == "preparing" && free && reachable { + preparing = true + } + if !n.Online || !n.ServingReady || n.ReadyGeneration == nil || !free || !reachable { + continue + } + if chosen == nil || *n.ReadyGeneration > *chosen.ReadyGeneration || *n.ReadyGeneration == *chosen.ReadyGeneration && n.Active < chosen.Active { + chosen = n + } + } + if chosen == nil { + if preparing { + return nil, ErrNodesPreparing + } + return nil, ErrNodeUnavailable + } + return &Placement{NodeID: chosen.ID, Generation: *chosen.ReadyGeneration}, nil +} + +// CheckReserved admits an allocation on the node placement its Session +// reserved: the placement must be unreleased and its node available. +func CheckReserved(reserved Reserved) error { + if reserved.Released || !reserved.Available { + return ErrNodeUnavailable + } + return nil +} + +// CheckRestore admits restoring suspended compute on its original node: the +// node must be online, ready for the allocation's generation and below its +// active capacity. Restore never selects another node. +func CheckRestore(restore Restore) error { + n := restore.Node + if n == nil || restore.Generation == 0 || !n.Online || !restore.GenerationReady || n.Active >= int64(n.MaxActive) { + return ErrNodeUnavailable + } + return nil +} + +// LoopbackOrigin reports whether a validated origin names a loopback host, +// which nothing outside the Core host can reach. +func LoopbackOrigin(value string) bool { + u, err := url.Parse(value) + if err != nil { + return false + } + if ip := net.ParseIP(u.Hostname()); ip != nil { + return ip.IsLoopback() + } + return u.Hostname() == "localhost" +} diff --git a/services/core/internal/deployment/placement/placement_test.go b/services/core/internal/deployment/placement/placement_test.go new file mode 100644 index 000000000..ff7212cde --- /dev/null +++ b/services/core/internal/deployment/placement/placement_test.go @@ -0,0 +1,204 @@ +package placement + +import ( + "encoding/json" + "errors" + "testing" + + "github.com/google/uuid" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" +) + +const publicURL = "https://core.example" + +// fakeDeclarations serves the declarations a test sets; a declaration the +// test left nil fails the test. +type fakeDeclarations struct { + t *testing.T + requiresPublicOrigin func(provider string) (bool, error) + validateSpecification func(provider string, spec sandbox.DeploymentSpec) error +} + +func (f *fakeDeclarations) RequiresPublicOrigin(provider string) (bool, error) { + if f.requiresPublicOrigin == nil { + f.t.Fatalf("unexpected RequiresPublicOrigin(%q)", provider) + } + return f.requiresPublicOrigin(provider) +} + +func (f *fakeDeclarations) ValidateSpecification(provider string, spec sandbox.DeploymentSpec) error { + if f.validateSpecification == nil { + f.t.Fatalf("unexpected ValidateSpecification(%q)", provider) + } + return f.validateSpecification(provider, spec) +} + +func rules(t *testing.T, declarations Declarations, url string) *Rules { + t.Helper() + r, err := NewRules(declarations, url) + if err != nil { + t.Fatal(err) + } + return r +} + +func generation(value uint64) *uint64 { return &value } + +func TestNewRulesRequiresDeclarations(t *testing.T) { + if _, err := NewRules(nil, publicURL); err == nil { + t.Fatal("NewRules accepted nil declarations") + } + if got := rules(t, &fakeDeclarations{t: t}, publicURL).PublicURL(); got != publicURL { + t.Fatalf("PublicURL = %q", got) + } +} + +// The built-in registry satisfies the declarations the rules read. +func TestRegistryDeclaresPublicOrigin(t *testing.T) { + loopback := rules(t, providers.Builtin(), "http://127.0.0.1:8091") + if err := loopback.CheckPublicOrigin("e2b"); !errors.Is(err, ErrPublicURLUnreachable) { + t.Fatalf("CheckPublicOrigin(e2b) on loopback = %v", err) + } + if err := loopback.CheckPublicOrigin("docker"); err != nil { + t.Fatalf("CheckPublicOrigin(docker) on loopback = %v", err) + } + if err := rules(t, providers.Builtin(), publicURL).CheckPublicOrigin("e2b"); err != nil { + t.Fatalf("CheckPublicOrigin(e2b) on a public URL = %v", err) + } + if err := loopback.CheckPublicOrigin("unknown"); err == nil { + t.Fatal("CheckPublicOrigin accepted an unknown provider") + } +} + +func TestCheckAdmission(t *testing.T) { + installation := uuid.NewString() + invalid := errors.New("invalid specification") + declarations := &fakeDeclarations{t: t, validateSpecification: func(provider string, spec sandbox.DeploymentSpec) error { + if spec.Resources.CPUs != 2 { + return invalid + } + return nil + }} + valid := Deployment{InstallationID: installation, Provider: "docker", Specification: json.RawMessage(`{"resources":{"cpus":2}}`)} + with := func(change func(*Deployment)) Deployment { + d := valid + change(&d) + return d + } + for name, test := range map[string]struct { + d Deployment + installation string + want error + message string + }{ + "unclaimed admits everything": {Deployment{Resetting: true, AdmissionPaused: true}, installation, nil, ""}, + "admitted": {valid, installation, nil, ""}, + "new Session": {valid, "", nil, ""}, + "reset": {with(func(d *Deployment) { d.Resetting, d.AdmissionPaused = true, true }), installation, ErrResetAdmission, "hosted admission is paused for a sandbox reset"}, + "paused": {with(func(d *Deployment) { d.AdmissionPaused = true }), installation, ErrAdmissionClosed, "environment is no longer available: sandbox creation is paused for provider maintenance"}, + "malformed specification": {with(func(d *Deployment) { d.Specification = json.RawMessage(`[`) }), installation, ErrAdmissionClosed, "environment is no longer available: sandbox creation requires a deployment specification"}, + "rejected specification": {with(func(d *Deployment) { d.Specification = json.RawMessage(`{"resources":{"cpus":3}}`) }), installation, ErrAdmissionClosed, "environment is no longer available: sandbox creation requires a deployment specification"}, + "no provider": {with(func(d *Deployment) { d.Provider, d.Specification = "", json.RawMessage(`[`) }), installation, nil, ""}, + "other installation": {valid, uuid.NewString(), ErrAdmissionClosed, "environment is no longer available: sandbox installation does not match deployment"}, + } { + err := rules(t, declarations, publicURL).CheckAdmission(test.d, test.installation) + if !errors.Is(err, test.want) || (test.want == nil) != (err == nil) || (err != nil && err.Error() != test.message) { + t.Errorf("%s: CheckAdmission = %v, want %v with %q", name, err, test.want, test.message) + } + } +} + +func TestDecidePlacement(t *testing.T) { + origin := func(required bool) *fakeDeclarations { + return &fakeDeclarations{t: t, requiresPublicOrigin: func(string) (bool, error) { return required, nil }} + } + ready := func(id string, g uint64, active int64) Node { + return Node{ID: id, Online: true, ServingReady: true, ReadyGeneration: generation(g), Active: active, MaxActive: 4, Retained: active, MaxRetained: 4, CoreURL: publicURL} + } + nodes := Deployment{Provider: "docker", Mode: "nodes", WebManaged: true} + with := func(change func(*Deployment)) Deployment { + d := nodes + change(&d) + return d + } + offline, unready, full, retainedFull, elsewhere := ready("offline", 9, 0), ready("unready", 9, 0), ready("full", 9, 4), ready("retained", 9, 0), ready("elsewhere", 9, 0) + offline.Online = false + unready.ServingReady = false + retainedFull.Retained = 4 + elsewhere.CoreURL = "https://other.example" + preparing := Node{ID: "preparing", Online: true, TargetState: "preparing", MaxActive: 1, MaxRetained: 1, CoreURL: publicURL} + for name, test := range map[string]struct { + declarations *fakeDeclarations + url string + d Deployment + nodes []Node + want *Placement + err error + }{ + "reset": {origin(false), publicURL, with(func(d *Deployment) { d.Resetting = true }), nil, nil, ErrResetAdmission}, + "loopback public origin": {origin(true), "http://localhost:8091", nodes, []Node{ready("a", 1, 0)}, nil, ErrPublicURLUnreachable}, + "direct": {origin(false), publicURL, with(func(d *Deployment) { d.Mode = "direct" }), nil, nil, nil}, + "direct paused": {origin(false), publicURL, with(func(d *Deployment) { d.Mode, d.AdmissionPaused = "direct", true }), nil, nil, ErrNodeUnavailable}, + "no provider": {&fakeDeclarations{t: t}, publicURL, Deployment{}, nil, nil, nil}, + "no provider on Web": {&fakeDeclarations{t: t}, publicURL, Deployment{WebManaged: true}, nil, nil, ErrNodeUnavailable}, + "paused": {origin(false), publicURL, with(func(d *Deployment) { d.AdmissionPaused = true }), []Node{ready("a", 1, 0)}, nil, ErrNodeUnavailable}, + "no nodes": {origin(false), publicURL, nodes, nil, nil, ErrNodeUnavailable}, + "only ineligible nodes": {origin(false), publicURL, nodes, []Node{offline, unready, full, retainedFull, elsewhere, {ID: "never", Online: true, ServingReady: true, MaxActive: 1, MaxRetained: 1, CoreURL: publicURL}}, nil, ErrNodeUnavailable}, + "preparing": {origin(false), publicURL, nodes, []Node{offline, preparing}, nil, ErrNodesPreparing}, + "highest generation": {origin(false), publicURL, nodes, []Node{ready("old", 1, 0), ready("new", 2, 3), preparing}, &Placement{NodeID: "new", Generation: 2}, nil}, + "fewest active": {origin(false), publicURL, nodes, []Node{ready("busy", 2, 3), ready("idle", 2, 1)}, &Placement{NodeID: "idle", Generation: 2}, nil}, + "other address off Web": {origin(false), publicURL, with(func(d *Deployment) { d.WebManaged = false }), []Node{elsewhere}, &Placement{NodeID: "elsewhere", Generation: 9}, nil}, + "public origin on public URL": {origin(true), publicURL, nodes, []Node{ready("a", 1, 0)}, &Placement{NodeID: "a", Generation: 1}, nil}, + } { + got, err := rules(t, test.declarations, test.url).DecidePlacement(test.d, test.nodes) + if !errors.Is(err, test.err) || (test.err == nil) != (err == nil) || (got == nil) != (test.want == nil) || (got != nil && *got != *test.want) { + t.Errorf("%s: DecidePlacement = %+v, %v; want %+v, %v", name, got, err, test.want, test.err) + } + } +} + +func TestCheckReservedAndRestore(t *testing.T) { + for name, test := range map[string]struct { + reserved Reserved + want error + }{ + "available": {Reserved{NodeID: "a", Generation: 1, Available: true}, nil}, + "released": {Reserved{NodeID: "a", Generation: 1, Released: true, Available: true}, ErrNodeUnavailable}, + "unavailable": {Reserved{NodeID: "a", Generation: 1}, ErrNodeUnavailable}, + } { + if err := CheckReserved(test.reserved); !errors.Is(err, test.want) || (test.want == nil) != (err == nil) { + t.Errorf("%s: CheckReserved = %v, want %v", name, err, test.want) + } + } + node := func(online bool, active int64) *Node { + return &Node{ID: "a", Online: online, Active: active, MaxActive: 2} + } + for name, test := range map[string]struct { + restore Restore + want error + }{ + "ready": {Restore{Node: node(true, 1), Generation: 3, GenerationReady: true}, nil}, + "missing node": {Restore{Generation: 3, GenerationReady: true}, ErrNodeUnavailable}, + "no generation": {Restore{Node: node(true, 1), GenerationReady: true}, ErrNodeUnavailable}, + "offline": {Restore{Node: node(false, 1), Generation: 3, GenerationReady: true}, ErrNodeUnavailable}, + "generation unready": {Restore{Node: node(true, 1), Generation: 3}, ErrNodeUnavailable}, + "at capacity": {Restore{Node: node(true, 2), Generation: 3, GenerationReady: true}, ErrNodeUnavailable}, + } { + if err := CheckRestore(test.restore); !errors.Is(err, test.want) || (test.want == nil) != (err == nil) { + t.Errorf("%s: CheckRestore = %v, want %v", name, err, test.want) + } + } +} + +func TestLoopbackOrigin(t *testing.T) { + for value, want := range map[string]bool{ + "http://localhost:8091": true, "http://127.0.0.1:8091": true, "http://127.0.0.2": true, "http://[::1]:8091": true, + "https://core.example": false, "https://[2001:db8::1]": false, "https://10.0.0.1": false, "": false, "http://host.localhost": false, + } { + if got := LoopbackOrigin(value); got != want { + t.Errorf("LoopbackOrigin(%q) = %v, want %v", value, got, want) + } + } +} diff --git a/services/core/internal/deployment/public_url.go b/services/core/internal/deployment/public_url.go index bfa5f5178..a676955f7 100644 --- a/services/core/internal/deployment/public_url.go +++ b/services/core/internal/deployment/public_url.go @@ -49,19 +49,6 @@ func ValidateCoreURL(value string) error { return nil } -// LoopbackOrigin reports whether a validated origin names a loopback host, which -// nothing outside the Core host can reach. -func LoopbackOrigin(value string) bool { - u, err := url.Parse(value) - if err != nil { - return false - } - if ip := net.ParseIP(u.Hostname()); ip != nil { - return ip.IsLoopback() - } - return u.Hostname() == "localhost" -} - // AddressBindings counts what is bound to an installation address: nodes // connect to the address they enrolled with, hosted sandboxes were started with // the address current at the time, and self-hosted executors were installed @@ -75,5 +62,5 @@ type AddressBindings struct { // AddressBindings counts what is bound to the installation public URL. func (s *Service) AddressBindings(ctx context.Context) (AddressBindings, error) { - return s.reader.AddressBindings(ctx, s.publicURL) + return s.reader.AddressBindings(ctx, s.rules.PublicURL()) } diff --git a/services/core/internal/deployment/rules_test.go b/services/core/internal/deployment/rules_test.go index 1812f531d..8cefc2f86 100644 --- a/services/core/internal/deployment/rules_test.go +++ b/services/core/internal/deployment/rules_test.go @@ -34,17 +34,6 @@ func TestValidateCoreURL(t *testing.T) { } } -func TestLoopbackOrigin(t *testing.T) { - for value, want := range map[string]bool{ - "http://localhost:8091": true, "http://127.0.0.1:8091": true, "http://127.0.0.2": true, "http://[::1]:8091": true, - "https://core.example": false, "https://[2001:db8::1]": false, "https://10.0.0.1": false, "": false, "http://host.localhost": false, - } { - if got := LoopbackOrigin(value); got != want { - t.Errorf("LoopbackOrigin(%q) = %v, want %v", value, got, want) - } - } -} - func TestParseID(t *testing.T) { id := uuid.New() if got, err := parseID(strings.ToUpper(id.String())); err != nil || got != id.String() { diff --git a/services/core/internal/deployment/service.go b/services/core/internal/deployment/service.go index 6be42ba1a..673a41dac 100644 --- a/services/core/internal/deployment/service.go +++ b/services/core/internal/deployment/service.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" ) @@ -15,18 +16,20 @@ type Service struct { storage Storage reader Reader registry *providers.Registry - // publicURL is OAC_PUBLIC_URL. Core reports it as the deployment and node - // configuration core_url and records it on each node it enrolls. - publicURL string + // rules decide admission and placement. Their public URL is OAC_PUBLIC_URL: + // Core reports it as the deployment and node configuration core_url and + // records it on each node it enrolls. + rules *placement.Rules } -// NewService returns the deployment service. publicURL is the validated -// installation public URL, empty when the installation has none. -func NewService(storage Storage, reader Reader, registry *providers.Registry, publicURL string) (*Service, error) { - if storage == nil || reader == nil || registry == nil { - return nil, errors.New("deployment service requires storage, a reader and a provider registry") - } - return &Service{storage: storage, reader: reader, registry: registry, publicURL: publicURL}, nil +// NewService returns the deployment service. rules are the installation's +// placement rules, built once with its provider declarations and validated +// public URL. +func NewService(storage Storage, reader Reader, registry *providers.Registry, rules *placement.Rules) (*Service, error) { + if storage == nil || reader == nil || registry == nil || rules == nil { + return nil, errors.New("deployment service requires storage, a reader, a provider registry and placement rules") + } + return &Service{storage: storage, reader: reader, registry: registry, rules: rules}, nil } // View returns the deployment as administrators read it. @@ -41,7 +44,7 @@ func (s *Service) View(ctx context.Context) (View, error) { // view reports the public URL as the deployment's read-only core_url. func (s *Service) view(snapshot Snapshot) (View, error) { d := snapshot.Record - result := View{InstallationID: d.InstallationID, Provider: d.Provider, CoreURL: s.publicURL, OwnerEpoch: d.OwnerEpoch, Generation: d.Generation, Mode: d.Mode, Rollout: snapshot.Rollout, Resources: snapshot.Resources} + result := View{InstallationID: d.InstallationID, Provider: d.Provider, CoreURL: s.rules.PublicURL(), OwnerEpoch: d.OwnerEpoch, Generation: d.Generation, Mode: d.Mode, Rollout: snapshot.Rollout, Resources: snapshot.Resources} if len(d.Specification) > 0 && string(d.Specification) != "{}" { var spec sandbox.DeploymentSpec if json.Unmarshal(d.Specification, &spec) == nil { @@ -247,13 +250,9 @@ func (s *Service) SetupForSelection(installationID string, input sandbox.Selecti return Setup{}, configurationError(err) } // Adapters declare whether their guests require a public Core origin. - publicOrigin, err := s.registry.RequiresPublicOrigin(input.Provider) - if err != nil { + if err := s.rules.CheckPublicOrigin(input.Provider); err != nil { return Setup{}, err } - if publicOrigin && LoopbackOrigin(s.publicURL) { - return Setup{}, ErrPublicURLUnreachable - } result := Setup{InstallationID: installationID, Provider: input.Provider, Mode: description.Mode, Specification: normalized.DeploymentSpec, Configuration: normalized.Configuration, BackendFingerprint: description.BackendFingerprint} return s.describe(result, description.IdleSeconds, description.RetentionSeconds) } diff --git a/services/core/internal/deployment/service_test.go b/services/core/internal/deployment/service_test.go index 92560479b..9fc5964aa 100644 --- a/services/core/internal/deployment/service_test.go +++ b/services/core/internal/deployment/service_test.go @@ -10,6 +10,7 @@ import ( "github.com/google/uuid" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/providercontract" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/e2b" @@ -20,13 +21,24 @@ const testPublicURL = "https://core.example" func newService(t *testing.T, storage *fakeStorage, reader *fakeReader, publicURL string) *Service { t.Helper() - service, err := NewService(storage, reader, providers.Builtin(), publicURL) + registry := providers.Builtin() + service, err := NewService(storage, reader, registry, newRules(t, registry, publicURL)) if err != nil { t.Fatal(err) } return service } +// newRules builds the placement rules as cmd/server does. +func newRules(t *testing.T, registry *providers.Registry, publicURL string) *placement.Rules { + t.Helper() + rules, err := placement.NewRules(registry, publicURL) + if err != nil { + t.Fatal(err) + } + return rules +} + // operations builds execution operations whose every transaction runs on tx. func operations(t *testing.T, publicURL string, tx *fakeDeploymentTx) *ExecutionOperations { t.Helper() @@ -65,10 +77,12 @@ func webDeployment(t *testing.T, installation, provider string, generation uint6 func TestNewServiceAndOperationsRejectNilDependencies(t *testing.T) { storage, reader, registry := &fakeStorage{t: t}, &fakeReader{t: t}, providers.Builtin() + rules := newRules(t, registry, testPublicURL) for name, build := range map[string]func() (*Service, error){ - "storage": func() (*Service, error) { return NewService(nil, reader, registry, testPublicURL) }, - "reader": func() (*Service, error) { return NewService(storage, nil, registry, testPublicURL) }, - "registry": func() (*Service, error) { return NewService(storage, reader, nil, testPublicURL) }, + "storage": func() (*Service, error) { return NewService(nil, reader, registry, rules) }, + "reader": func() (*Service, error) { return NewService(storage, nil, registry, rules) }, + "registry": func() (*Service, error) { return NewService(storage, reader, nil, rules) }, + "rules": func() (*Service, error) { return NewService(storage, reader, registry, nil) }, } { if service, err := build(); service != nil || err == nil { t.Errorf("NewService without %s = %v, %v", name, service, err) @@ -89,7 +103,7 @@ func TestSetupForSelectionRequiresAReachablePublicURL(t *testing.T) { id := uuid.NewString() e2bSelection := sandbox.Selection{Provider: "e2b", Configuration: &e2b.DeploymentConfiguration{APIKey: "synthetic-key", Template: "runtime:" + uuid.NewString()}} for _, publicURL := range []string{"http://127.0.0.1:8091", "http://localhost:8091", "http://[::1]:8091"} { - if _, err := newService(t, &fakeStorage{t: t}, &fakeReader{t: t}, publicURL).SetupForSelection(id, e2bSelection); !errors.Is(err, ErrPublicURLUnreachable) { + if _, err := newService(t, &fakeStorage{t: t}, &fakeReader{t: t}, publicURL).SetupForSelection(id, e2bSelection); !errors.Is(err, placement.ErrPublicURLUnreachable) { t.Errorf("E2B accepted the loopback public URL %s: %v", publicURL, err) } } @@ -291,8 +305,8 @@ func TestEnrollChecksTheTokenBeforeTheDeployment(t *testing.T) { {"consumed token before setup", Record{}, consumed, nil, valid, ErrNodeCredential}, {"expired token while paused", paused, expired, nil, valid, ErrNodeCredential}, {"token of another installation while resetting", resetting, foreign, nil, valid, ErrNodeCredential}, - {"token store failure", current, EnrollmentRecord{}, errors.New("database down"), valid, ErrNodeUnavailable}, - {"no provider selected", unselected, receipt, nil, valid, ErrNodeUnavailable}, + {"token store failure", current, EnrollmentRecord{}, errors.New("database down"), valid, placement.ErrNodeUnavailable}, + {"no provider selected", unselected, receipt, nil, valid, placement.ErrNodeUnavailable}, {"reset in progress", resetting, receipt, nil, valid, ErrResetInProgress}, {"admission paused", paused, receipt, nil, valid, ErrInvalidInput}, {"stale generation", current, receipt, nil, stale, ErrSpecificationMismatch}, @@ -349,7 +363,7 @@ func TestAuthenticateNodeChecksTheCredentialFirst(t *testing.T) { want error }{ {"unknown node", StoredNode{}, ErrNotFound, credential, ErrNodeCredential}, - {"node store failure", StoredNode{}, errors.New("database down"), credential, ErrNodeUnavailable}, + {"node store failure", StoredNode{}, errors.New("database down"), credential, placement.ErrNodeUnavailable}, {"wrong credential", stored, nil, strings.Repeat("x", 64), ErrNodeCredential}, } { reads := &fakeNodeReads{t: t, loadNode: func(id string) (StoredNode, error) { diff --git a/services/core/internal/deployment/storage.go b/services/core/internal/deployment/storage.go index 05f41e34f..399c98d88 100644 --- a/services/core/internal/deployment/storage.go +++ b/services/core/internal/deployment/storage.go @@ -7,7 +7,9 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/adminaudit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/coremetrics" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) // Storage is the persistence node management writes through, on pooled @@ -28,6 +30,17 @@ type Storage interface { // host history, at the Runtime sampler cadence, and returns the number of // samples. Neither reads nor offline nodes fill gaps. SampleHostHistory(ctx context.Context) (int64, error) + // WithActivity runs apply in one transaction that locks the Session of + // the tenant's Environment and then prunes the Session's journal. It + // commits only when both succeed. A missing Environment, or one of a + // deleted Session, is sessions.ErrNotFound. + WithActivity(ctx context.Context, key AllocationKey, apply func(sessions.LockedSession, ActivityTx) error) error +} + +// ActivityTx is one Session-locked activity change. +type ActivityTx interface { + // TouchActivity records live-compute activity of the Environment. + TouchActivity() error } // ExecutionStorage is the persistence deployment changes write through. Only @@ -37,6 +50,87 @@ type ExecutionStorage interface { // WithDeployment runs apply in one leased transaction that begins by // locking the deployment. It commits only when apply returns nil. WithDeployment(ctx context.Context, apply func(DeploymentTx) error) error + // WithReservation runs apply in one leased transaction that locks the + // Session of the tenant's Environment and then prunes the Session's + // journal. It commits only when both succeed. A missing Environment, or + // one of a deleted Session, is sessions.ErrNotFound. + WithReservation(ctx context.Context, key AllocationKey, apply func(sessions.LockedSession, ReservationTx) error) error + // WithAllocation runs apply in one leased transaction that locks the + // Session owning the Environment's allocation, deleted or not, and then + // prunes the Session's journal. It commits only when both succeed. A + // missing allocation is ErrNotFound. + WithAllocation(ctx context.Context, key AllocationKey, apply func(AllocationTx) error) error + AllocationCleanupStorage + // ClearWake clears the allocation's wake request unless activity newer + // than observed arrived. It runs on the lease. + ClearWake(ctx context.Context, allocationID string, observed time.Time) error +} + +// AllocationCleanupStorage runs allocation cleanup, which settles the +// owning Session's work and Environment with the allocation. +type AllocationCleanupStorage interface { + // WithAllocationCleanup runs apply as WithAllocation does, with the + // Session's cancellation and Environment termination bound to the same + // transaction. + WithAllocationCleanup(ctx context.Context, key AllocationKey, apply func(AllocationCleanupTx) error) error +} + +// ReservationTx is one Session-locked allocation reservation. The Session is +// the owner of the Environment the transaction was opened for. +type ReservationTx interface { + sessions.EnvironmentDeviceTx + // LoadEnvironment reads the Session's Environment. + LoadEnvironment(ctx context.Context) (sessions.Environment, error) + // FindAllocation returns the Environment's allocation and whether it has + // one. + FindAllocation() (Allocation, bool, error) + // LockDeployment locks the deployment and returns it as placement reads it. + LockDeployment() (placement.Deployment, error) + // LoadReserved returns the node placement the Session reserved for the + // Environment. + LoadReserved() (placement.Reserved, error) + // InsertAllocation stores the allocation of the Session's Environment. + InsertAllocation(allocation NewAllocation) (Allocation, error) +} + +// AllocationTx is one Session-locked allocation change. Each write applies to +// current, the allocation LoadAllocation returned, and returns it changed. A +// write whose stored guard no longer holds is ErrAllocationConflict. +type AllocationTx interface { + // LoadAllocation returns the Environment's allocation. + LoadAllocation() (Allocation, error) + // LoadSessionDevice returns the device the Session is bound to and + // whether it is bound to one. + LoadSessionDevice() (SessionDevice, bool, error) + // LoadActivity returns the allocation's activity. + LoadActivity(current Allocation) (Activity, error) + // LoadRestore locks the deployment and returns what restoring the + // allocation's suspended compute on its node reads. + LoadRestore(current Allocation) (placement.Restore, error) + // ObserveRunning records the allocation running with its creation settled. + ObserveRunning(current Allocation) (Allocation, error) + // Keep renews the allocation's lease. + Keep(current Allocation) (Allocation, error) + // SettleCreation records that the original Create can no longer change + // resources. + SettleCreation(current Allocation) (Allocation, error) + // Release releases the allocation and its node placement. + Release(current Allocation) (Allocation, error) + // SetCompute commits the compute phase and advances its revision. + SetCompute(current Allocation, change ComputeChange) (Allocation, error) + // RecordObservation records the diagnostic for current's compute + // revision and state. It changes nothing once either moved on. + RecordObservation(current Allocation, diagnostic string) error +} + +// AllocationCleanupTx is one Session-locked allocation cleanup. +type AllocationCleanupTx interface { + AllocationTx + sessions.EnvironmentTerminationTx + // RevokeDevice revokes the allocation's device. + RevokeDevice(current Allocation) error + // RequestCleanup records that the allocation's resources await cleanup. + RequestCleanup(current Allocation) (Allocation, error) } // Reader answers deployment and node queries. @@ -72,6 +166,49 @@ type Reader interface { // AddressBindings counts what is bound to an installation address, read // in one snapshot. publicURL is the address nodes are compared against. AddressBindings(ctx context.Context, publicURL string) (AddressBindings, error) + // EnvironmentAllocation returns the Environment's allocation, including + // for a deleted Session. It returns ErrInvalidInput for a malformed + // identifier and ErrNotFound for a missing allocation. + EnvironmentAllocation(ctx context.Context, key AllocationKey) (Allocation, error) + // CredentialAllocations returns up to 32 unreleased allocations after the + // given allocation ID, in ID order. An empty after starts at the + // beginning. + CredentialAllocations(ctx context.Context, after string) ([]Allocation, error) + // ObservationSessions returns up to limit undeleted hosted Sessions after + // the given Session ID, in ID order, whose allocation, if any, is + // unreleased. limit is 1 to 100. It reads only and renews nothing. + ObservationSessions(ctx context.Context, after string, limit int) (ObservationSessionPage, error) + // NodeAllocations returns up to 1000 unreleased allocations of the node, + // oldest first. A missing node is ErrNotFound. + NodeAllocations(ctx context.Context, nodeID string) ([]NodeAllocation, error) + // NodeOnline reports whether the node is connected; a missing node is + // offline. + NodeOnline(ctx context.Context, nodeID string) (bool, error) + // LifecycleNodes returns the node of each allocation lifecycle, offline + // ones included, with the empty ID for the lifecycle without a node. + LifecycleNodes(ctx context.Context) ([]string, error) + // LifecycleAllocations returns up to 32 allocations of the node's + // lifecycle after the given allocation ID, in ID order. An empty nodeID + // names the lifecycle without a node. + LifecycleAllocations(ctx context.Context, nodeID, after string) ([]Allocation, error) + // UnallocatedEnvironments returns up to 32 pending hosted Environments of + // the node's lifecycle without an allocation, after the given Environment + // ID, in ID order. It reads the committed placement and never selects a + // replacement node. + UnallocatedEnvironments(ctx context.Context, nodeID, after string) ([]UnallocatedEnvironment, error) + // LifecyclePlacement returns what routing the Environment to its + // lifecycle reads, including for a deleted Session. A missing Environment + // is sessions.ErrNotFound. + LifecyclePlacement(ctx context.Context, key AllocationKey) (LifecyclePlacement, error) + // Activity returns the allocation's activity; a missing allocation is + // ErrNotFound. + Activity(ctx context.Context, allocationID string) (Activity, error) + // CountComputeReservations counts the installation's allocations that + // reserve active compute. + CountComputeReservations(ctx context.Context, installationID string) (int64, error) + // CountRetainedAllocations counts the installation's allocations that + // retain resources. + CountRetainedAllocations(ctx context.Context, installationID string) (int64, error) } // NodeReads loads node authentication facts. diff --git a/services/core/internal/execution/archive_cancellation_cleanup_test.go b/services/core/internal/execution/archive_cancellation_cleanup_test.go index d3f29bbde..6618cb37f 100644 --- a/services/core/internal/execution/archive_cancellation_cleanup_test.go +++ b/services/core/internal/execution/archive_cancellation_cleanup_test.go @@ -9,6 +9,7 @@ import ( "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/projectpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" @@ -64,7 +65,7 @@ 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, leased, _, _, pool := resetManagerStoreDB(t, nil) + s, leased, deployments, reader, pool := resetManagerStoreDB(t, nil) writer := leased.Store installation := initializeE2BDeployment(t, leased) projectID := uuid.NewString() @@ -82,11 +83,12 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { t.Fatal(err) } secret := uuid.NewString() - owner, err := writer.ReserveRuntimeAllocation(t.Context(), project.TenantID, session.Environment.ID, installation, runtimedevice.HashCredential(secret)) + key := deployment.AllocationKey{TenantID: project.TenantID, EnvironmentID: session.Environment.ID} + owner, err := leased.Deployment.ReserveAllocation(t.Context(), key, installation, runtimedevice.HashCredential(secret)) if err != nil { t.Fatal(err) } - owner, err = writer.ObserveRuntimeRunning(t.Context(), owner) + owner, err = leased.Deployment.ObserveRunning(t.Context(), owner) if err != nil { t.Fatal(err) } @@ -104,7 +106,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { currentCompute := sandbox.Compute{ID: uuid.NewString(), Name: owner.ID + "-g0"} if checkpoint { state, _ := json.Marshal(runtimeCompute{Current: currentCompute}) - owner, err = writer.SetRuntimeCompute(t.Context(), owner, "running", state, nil, 0) + owner, err = leased.Deployment.SetCompute(t.Context(), owner, "running", state, nil, 0) if err != nil { t.Fatal(err) } @@ -149,7 +151,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { if _, err := writer.ArchiveManagedSession(audit, project.TenantID, session.ID, 1); err != nil { t.Fatal(err) } - owner, err = s.GetRuntimeAllocation(t.Context(), project.TenantID, session.Environment.ID) + owner, err = reader.EnvironmentAllocation(t.Context(), key) if err != nil { t.Fatal(err) } @@ -161,13 +163,13 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { if err != nil || turn.Status != expectedStatus || turn.CancelRequestedAt.IsZero() || (turn.CompletedAt.IsZero() != (expectedStatus == sessions.TurnWaiting)) { t.Fatal("cleanup observed unexpected terminal state", turn, err) } - allocation, err := s.GetRuntimeAllocation(t.Context(), project.TenantID, session.Environment.ID) + allocation, err := reader.EnvironmentAllocation(t.Context(), key) if err != nil || allocation.State != "cleanup_pending" { t.Fatal("Kill bypassed durable cleanup ownership", allocation, err) } }} - reader, _ := testSessions(t, pool, testCredentialCipher(t)) - lifecycle := &runtimeLifecycle{store: writer, sessions: reader, sessionExecution: leased.Sessions, lease: leased.Lease, registry: registry, config: RuntimeProvider{InstallationID: installation, Provider: provider}, connections: map[string]*runtimeConnection{}} + sessionReader, _ := testSessions(t, pool, testCredentialCipher(t)) + lifecycle := &runtimeLifecycle{store: writer, sessions: sessionReader, sessionExecution: leased.Sessions, deployment: leased.Deployment, deployments: deployments, reader: reader, lease: leased.Lease, registry: registry, config: RuntimeProvider{InstallationID: installation, Provider: provider}, connections: map[string]*runtimeConnection{}} if checkpoint { lifecycle.config.Provider = waitingCleanupCheckpoint{beforeKill: provider.beforeKill} } @@ -178,7 +180,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { if kills != 0 { t.Fatal("destroyed compute before cancellation committed") } - pending, err := s.GetRuntimeAllocation(t.Context(), project.TenantID, session.Environment.ID) + pending, err := reader.EnvironmentAllocation(t.Context(), key) if err != nil || pending.State != "cleanup_pending" { t.Fatal(pending, err) } @@ -191,7 +193,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { t.Fatal(err) } } - after, err := s.GetRuntimeAllocation(t.Context(), project.TenantID, session.Environment.ID) + after, err := reader.EnvironmentAllocation(t.Context(), key) if err != nil || after.State != "released" || kills != 1 { t.Fatal("cleanup did not release", after, kills, err) } diff --git a/services/core/internal/execution/deployment_fixture_test.go b/services/core/internal/execution/deployment_fixture_test.go index a1c794fa6..2c21c8010 100644 --- a/services/core/internal/execution/deployment_fixture_test.go +++ b/services/core/internal/execution/deployment_fixture_test.go @@ -4,22 +4,35 @@ import ( "context" "errors" "testing" + "time" "github.com/jackc/pgx/v5/pgxpool" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/coremetrics" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) // fixturePublicURL is the installation public URL of these tests. E2B requires // a public origin, so it is never loopback. const fixturePublicURL = "https://core.example" +// fixtureRules are the placement rules of these tests, on fixturePublicURL. +func fixtureRules(t *testing.T) *placement.Rules { + t.Helper() + rules, err := placement.NewRules(providers.Builtin(), fixturePublicURL) + if err != nil { + t.Fatal(err) + } + return rules +} + // testDeployment builds the pooled deployment service and reader and the // deployment execution operations on lease, as cmd/server does for the Worker. // cipher is nil when the owner has no credential key. @@ -43,7 +56,7 @@ func unitDeploymentService(t *testing.T) *deployment.Service { // the execution operations on execution. func deploymentOperations(t *testing.T, storage deployment.Storage, reader deployment.Reader, execution deployment.ExecutionStorage) (*deployment.Service, *deployment.ExecutionOperations) { t.Helper() - service, err := deployment.NewService(storage, reader, providers.Builtin(), fixturePublicURL) + service, err := deployment.NewService(storage, reader, providers.Builtin(), fixtureRules(t)) if err != nil { t.Fatal(err) } @@ -68,6 +81,7 @@ type strictDeploymentStorage struct { connectNode func(context.Context, string, string, uint64) (bool, error) disconnectNode func(context.Context, string, string, uint64) error sampleHostHistory func(context.Context) (int64, error) + withActivity func(context.Context, deployment.AllocationKey, func(sessions.LockedSession, deployment.ActivityTx) error) error } func (s *strictDeploymentStorage) WithNodes(ctx context.Context, apply func(deployment.NodeTx) error) error { @@ -101,8 +115,12 @@ func (s *strictDeploymentStorage) SampleHostHistory(ctx context.Context) (int64, // strictExecutionStorage runs withDeployment when set; otherwise any call // fails the test. type strictExecutionStorage struct { - t *testing.T - withDeployment func(context.Context, func(deployment.DeploymentTx) error) error + t *testing.T + withDeployment func(context.Context, func(deployment.DeploymentTx) error) error + withReservation func(context.Context, deployment.AllocationKey, func(sessions.LockedSession, deployment.ReservationTx) error) error + withAllocation func(context.Context, deployment.AllocationKey, func(deployment.AllocationTx) error) error + withAllocationCleanup func(context.Context, deployment.AllocationKey, func(deployment.AllocationCleanupTx) error) error + clearWake func(context.Context, string, time.Time) error } func (s *strictExecutionStorage) WithDeployment(ctx context.Context, apply func(deployment.DeploymentTx) error) error { @@ -114,17 +132,29 @@ func (s *strictExecutionStorage) WithDeployment(ctx context.Context, apply func( // strictDeploymentReader runs each set func; any other call fails the test. type strictDeploymentReader struct { - t *testing.T - deployment func(context.Context) (deployment.Record, error) - snapshot func(context.Context) (deployment.Snapshot, error) - ownerEpoch func(context.Context) (uint64, error) - allocation func(context.Context, sandbox.Reference) (deployment.AllocationRecord, error) - generations func(context.Context, int64) ([]deployment.GenerationRecord, error) - nodes func(context.Context) ([]deployment.NodeRecord, error) - nodeHistory func(context.Context, string, coremetrics.Range) (deployment.NodeRecord, []deployment.HostHistoryPoint, error) - readNodes func(context.Context, func(deployment.NodeReads) error) error - resetSessions func(context.Context, string, bool) ([]deployment.ResetSession, error) - addressBindings func(context.Context, string) (deployment.AddressBindings, error) + t *testing.T + deployment func(context.Context) (deployment.Record, error) + snapshot func(context.Context) (deployment.Snapshot, error) + ownerEpoch func(context.Context) (uint64, error) + allocation func(context.Context, sandbox.Reference) (deployment.AllocationRecord, error) + generations func(context.Context, int64) ([]deployment.GenerationRecord, error) + nodes func(context.Context) ([]deployment.NodeRecord, error) + nodeHistory func(context.Context, string, coremetrics.Range) (deployment.NodeRecord, []deployment.HostHistoryPoint, error) + readNodes func(context.Context, func(deployment.NodeReads) error) error + resetSessions func(context.Context, string, bool) ([]deployment.ResetSession, error) + addressBindings func(context.Context, string) (deployment.AddressBindings, error) + environmentAllocation func(context.Context, deployment.AllocationKey) (deployment.Allocation, error) + credentialAllocations func(context.Context, string) ([]deployment.Allocation, error) + observationSessions func(context.Context, string, int) (deployment.ObservationSessionPage, error) + nodeAllocations func(context.Context, string) ([]deployment.NodeAllocation, error) + nodeOnline func(context.Context, string) (bool, error) + lifecycleNodes func(context.Context) ([]string, error) + lifecycleAllocations func(context.Context, string, string) ([]deployment.Allocation, error) + unallocatedEnvironments func(context.Context, string, string) ([]deployment.UnallocatedEnvironment, error) + lifecyclePlacement func(context.Context, deployment.AllocationKey) (deployment.LifecyclePlacement, error) + activity func(context.Context, string) (deployment.Activity, error) + countComputeReservations func(context.Context, string) (int64, error) + countRetainedAllocations func(context.Context, string) (int64, error) } func (r *strictDeploymentReader) Deployment(ctx context.Context) (deployment.Record, error) { @@ -196,3 +226,122 @@ func (r *strictDeploymentReader) AddressBindings(ctx context.Context, publicURL } return r.addressBindings(ctx, publicURL) } + +func (s *strictDeploymentStorage) WithActivity(ctx context.Context, key deployment.AllocationKey, apply func(sessions.LockedSession, deployment.ActivityTx) error) error { + if s.withActivity == nil { + return unexpectedDeploymentCall(s.t, "WithActivity") + } + return s.withActivity(ctx, key, apply) +} + +func (s *strictExecutionStorage) WithReservation(ctx context.Context, key deployment.AllocationKey, apply func(sessions.LockedSession, deployment.ReservationTx) error) error { + if s.withReservation == nil { + return unexpectedDeploymentCall(s.t, "WithReservation") + } + return s.withReservation(ctx, key, apply) +} + +func (s *strictExecutionStorage) WithAllocation(ctx context.Context, key deployment.AllocationKey, apply func(deployment.AllocationTx) error) error { + if s.withAllocation == nil { + return unexpectedDeploymentCall(s.t, "WithAllocation") + } + return s.withAllocation(ctx, key, apply) +} + +func (s *strictExecutionStorage) WithAllocationCleanup(ctx context.Context, key deployment.AllocationKey, apply func(deployment.AllocationCleanupTx) error) error { + if s.withAllocationCleanup == nil { + return unexpectedDeploymentCall(s.t, "WithAllocationCleanup") + } + return s.withAllocationCleanup(ctx, key, apply) +} + +func (s *strictExecutionStorage) ClearWake(ctx context.Context, allocationID string, observed time.Time) error { + if s.clearWake == nil { + return unexpectedDeploymentCall(s.t, "ClearWake") + } + return s.clearWake(ctx, allocationID, observed) +} + +func (r *strictDeploymentReader) EnvironmentAllocation(ctx context.Context, key deployment.AllocationKey) (deployment.Allocation, error) { + if r.environmentAllocation == nil { + return deployment.Allocation{}, unexpectedDeploymentCall(r.t, "EnvironmentAllocation") + } + return r.environmentAllocation(ctx, key) +} + +func (r *strictDeploymentReader) CredentialAllocations(ctx context.Context, after string) ([]deployment.Allocation, error) { + if r.credentialAllocations == nil { + return nil, unexpectedDeploymentCall(r.t, "CredentialAllocations") + } + return r.credentialAllocations(ctx, after) +} + +func (r *strictDeploymentReader) ObservationSessions(ctx context.Context, after string, limit int) (deployment.ObservationSessionPage, error) { + if r.observationSessions == nil { + return deployment.ObservationSessionPage{}, unexpectedDeploymentCall(r.t, "ObservationSessions") + } + return r.observationSessions(ctx, after, limit) +} + +func (r *strictDeploymentReader) NodeAllocations(ctx context.Context, nodeID string) ([]deployment.NodeAllocation, error) { + if r.nodeAllocations == nil { + return nil, unexpectedDeploymentCall(r.t, "NodeAllocations") + } + return r.nodeAllocations(ctx, nodeID) +} + +func (r *strictDeploymentReader) NodeOnline(ctx context.Context, nodeID string) (bool, error) { + if r.nodeOnline == nil { + return false, unexpectedDeploymentCall(r.t, "NodeOnline") + } + return r.nodeOnline(ctx, nodeID) +} + +func (r *strictDeploymentReader) LifecycleNodes(ctx context.Context) ([]string, error) { + if r.lifecycleNodes == nil { + return nil, unexpectedDeploymentCall(r.t, "LifecycleNodes") + } + return r.lifecycleNodes(ctx) +} + +func (r *strictDeploymentReader) LifecycleAllocations(ctx context.Context, nodeID string, after string) ([]deployment.Allocation, error) { + if r.lifecycleAllocations == nil { + return nil, unexpectedDeploymentCall(r.t, "LifecycleAllocations") + } + return r.lifecycleAllocations(ctx, nodeID, after) +} + +func (r *strictDeploymentReader) UnallocatedEnvironments(ctx context.Context, nodeID string, after string) ([]deployment.UnallocatedEnvironment, error) { + if r.unallocatedEnvironments == nil { + return nil, unexpectedDeploymentCall(r.t, "UnallocatedEnvironments") + } + return r.unallocatedEnvironments(ctx, nodeID, after) +} + +func (r *strictDeploymentReader) LifecyclePlacement(ctx context.Context, key deployment.AllocationKey) (deployment.LifecyclePlacement, error) { + if r.lifecyclePlacement == nil { + return deployment.LifecyclePlacement{}, unexpectedDeploymentCall(r.t, "LifecyclePlacement") + } + return r.lifecyclePlacement(ctx, key) +} + +func (r *strictDeploymentReader) Activity(ctx context.Context, allocationID string) (deployment.Activity, error) { + if r.activity == nil { + return deployment.Activity{}, unexpectedDeploymentCall(r.t, "Activity") + } + return r.activity(ctx, allocationID) +} + +func (r *strictDeploymentReader) CountComputeReservations(ctx context.Context, installationID string) (int64, error) { + if r.countComputeReservations == nil { + return 0, unexpectedDeploymentCall(r.t, "CountComputeReservations") + } + return r.countComputeReservations(ctx, installationID) +} + +func (r *strictDeploymentReader) CountRetainedAllocations(ctx context.Context, installationID string) (int64, error) { + if r.countRetainedAllocations == nil { + return 0, unexpectedDeploymentCall(r.t, "CountRetainedAllocations") + } + return r.countRetainedAllocations(ctx, installationID) +} diff --git a/services/core/internal/execution/environment_directory.go b/services/core/internal/execution/environment_directory.go index 3093a23e0..bba21b5dd 100644 --- a/services/core/internal/execution/environment_directory.go +++ b/services/core/internal/execution/environment_directory.go @@ -63,7 +63,7 @@ func (w *Worker) runDirectoryRead(owner context.Context, request directoryReadRe } touch, stop := context.WithTimeout(context.WithoutCancel(owner), 5*time.Second) defer stop() - if err := w.dispatcher.Store.TouchRuntimeActivity(touch, request.environment.TenantID, request.environment.ID); err != nil { + if err := w.dispatcher.Deployment.TouchActivity(touch, request.environment.TenantID, request.environment.ID); err != nil { result.err = err } }() diff --git a/services/core/internal/execution/runtime_compute.go b/services/core/internal/execution/runtime_compute.go index 835b935ec..7d08a1b33 100644 --- a/services/core/internal/execution/runtime_compute.go +++ b/services/core/internal/execution/runtime_compute.go @@ -7,9 +7,8 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -39,7 +38,7 @@ func (r *runtimeLifecycle) computeCapacity(ctx context.Context, key string) erro if policy == nil { return nil } - count, err := r.store.CountRuntimeComputeReservations(ctx, key) + count, err := r.reader.CountComputeReservations(ctx, key) if err != nil { return err } @@ -48,7 +47,7 @@ func (r *runtimeLifecycle) computeCapacity(ctx context.Context, key string) erro } return nil } -func (r *runtimeLifecycle) saveCompute(ctx context.Context, owner store.RuntimeAllocation, phase string, state runtimeCompute, until *time.Time) (store.RuntimeAllocation, error) { +func (r *runtimeLifecycle) saveCompute(ctx context.Context, owner deployment.Allocation, phase string, state runtimeCompute, until *time.Time) (deployment.Allocation, error) { raw, err := json.Marshal(state) if err != nil { return owner, err @@ -61,9 +60,9 @@ func (r *runtimeLifecycle) saveCompute(ctx context.Context, owner store.RuntimeA } idleTimeout = policy.IdleTimeout } - return r.store.SetRuntimeCompute(ctx, owner, phase, raw, until, idleTimeout) + return r.deployment.SetCompute(ctx, owner, phase, raw, until, idleTimeout) } -func (r *runtimeLifecycle) enableCompute(ctx context.Context, owner store.RuntimeAllocation) error { +func (r *runtimeLifecycle) enableCompute(ctx context.Context, owner deployment.Allocation) error { p, capabilityErr := sandbox.Checkpoint(r.config.Provider) if capabilityErr != nil { return capabilityErr @@ -83,7 +82,7 @@ func (r *runtimeLifecycle) enableCompute(ctx context.Context, owner store.Runtim return err } -func (r *runtimeLifecycle) observeCompute(ctx context.Context, owner store.RuntimeAllocation) error { +func (r *runtimeLifecycle) observeCompute(ctx context.Context, owner deployment.Allocation) error { p, capabilityErr := sandbox.Checkpoint(r.config.Provider) if capabilityErr != nil { return capabilityErr @@ -123,7 +122,7 @@ func (r *runtimeLifecycle) observeCompute(ctx context.Context, owner store.Runti } } -func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.CheckpointProvider, owner store.RuntimeAllocation, state runtimeCompute) error { +func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.CheckpointProvider, owner deployment.Allocation, state runtimeCompute) error { compute, err := p.GetCompute(ctx, runtimeReference(owner), state.Current) if err != nil { return err @@ -135,16 +134,16 @@ func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.Checkpoint if err != nil { return err } - if _, err := r.store.KeepRuntimeAllocation(ctx, owner); err != nil { + if _, err := r.deployment.KeepAllocation(ctx, owner); err != nil { return err } - activity, err := r.store.RuntimeActivity(ctx, owner) + activity, err := r.reader.Activity(ctx, owner.ID) if err != nil { return err } if activity.WakeRequested { // Clearing only the observed timestamp cannot consume a newer live request. - return r.store.ClearRuntimeWake(ctx, owner, owner.ComputeActivityAt) + return r.deployment.ClearWake(ctx, owner, owner.ComputeActivityAt) } policy := r.config.Suspension if policy == nil || !activity.ReadyToSuspend(policy.IdleTimeout) { @@ -163,7 +162,7 @@ func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.Checkpoint if !result.Accepted { // A rejected request did not park the daemon. Reset its idle clock so a // continuing file operation is not immediately interrupted by another try. - if err := r.store.TouchRuntimeActivity(ctx, owner.TenantID, owner.EnvironmentID); err != nil { + if err := r.deployments.TouchActivity(ctx, owner.TenantID, owner.EnvironmentID); err != nil { return err } _, err = r.saveCompute(ctx, next, "running", runtimeCompute{Current: state.Current}, nil) @@ -176,7 +175,7 @@ func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.Checkpoint // The Session-locked phase commit checks pending work and wake requests. // A competing request keeps its queue position and resumes this source. suspending, err := r.saveCompute(ctx, next, "suspending", state, &until) - if errors.Is(err, sessions.ErrTurnConflict) { + if errors.Is(err, deployment.ErrAllocationConflict) { state.Rollback = true next, err = r.saveCompute(ctx, next, "waking", state, &until) if err != nil { @@ -190,7 +189,7 @@ func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.Checkpoint return r.captureCompute(ctx, p, suspending, state, false) } -func (r *runtimeLifecycle) captureCompute(ctx context.Context, p sandbox.CheckpointProvider, owner store.RuntimeAllocation, state runtimeCompute, observeOnly bool) error { +func (r *runtimeLifecycle) captureCompute(ctx context.Context, p sandbox.CheckpointProvider, owner deployment.Allocation, state runtimeCompute, observeOnly bool) error { result, err := p.Suspend(ctx, sandbox.SuspendRequest{Reference: runtimeReference(owner), OperationID: state.SuspendID, Source: state.Current, Snapshot: state.Snapshot, ObserveOnly: observeOnly}) if err != nil { return err @@ -223,8 +222,8 @@ func (r *runtimeLifecycle) captureCompute(ctx context.Context, p sandbox.Checkpo return err } -func (r *runtimeLifecycle) restoreIdleCompute(ctx context.Context, p sandbox.CheckpointProvider, owner store.RuntimeAllocation, state runtimeCompute) error { - activity, err := r.store.RuntimeActivity(ctx, owner) +func (r *runtimeLifecycle) restoreIdleCompute(ctx context.Context, p sandbox.CheckpointProvider, owner deployment.Allocation, state runtimeCompute) error { + activity, err := r.reader.Activity(ctx, owner.ID) if err != nil { return err } @@ -248,7 +247,7 @@ func (r *runtimeLifecycle) restoreIdleCompute(ctx context.Context, p sandbox.Che } return r.restoreCompute(ctx, p, next, state, false) } -func (r *runtimeLifecycle) restoreCompute(ctx context.Context, p sandbox.CheckpointProvider, owner store.RuntimeAllocation, state runtimeCompute, observeOnly bool) error { +func (r *runtimeLifecycle) restoreCompute(ctx context.Context, p sandbox.CheckpointProvider, owner deployment.Allocation, state runtimeCompute, observeOnly bool) error { if state.Target == nil || state.Snapshot == nil || state.Rollback { return sandbox.ErrOwnership } @@ -276,8 +275,8 @@ func ignoreComputeAbsent(err error) error { return err } -// Node-backed restores reserve capacity atomically in SetRuntimeCompute. -func (r *runtimeLifecycle) computeCapacityForAllocation(ctx context.Context, owner store.RuntimeAllocation) error { +// Node-backed restores reserve capacity atomically in SetCompute. +func (r *runtimeLifecycle) computeCapacityForAllocation(ctx context.Context, owner deployment.Allocation) error { if owner.NodeID != "" { return nil } diff --git a/services/core/internal/execution/runtime_compute_wake.go b/services/core/internal/execution/runtime_compute_wake.go index 339f7847c..552ede1ed 100644 --- a/services/core/internal/execution/runtime_compute_wake.go +++ b/services/core/internal/execution/runtime_compute_wake.go @@ -6,13 +6,13 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "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/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) -func (r *runtimeLifecycle) wakeCompute(ctx context.Context, p sandbox.CheckpointProvider, owner store.RuntimeAllocation, state runtimeCompute) error { +func (r *runtimeLifecycle) wakeCompute(ctx context.Context, p sandbox.CheckpointProvider, owner deployment.Allocation, state runtimeCompute) error { if state.Rollback { if _, err := p.ResumeCompute(ctx, runtimeReference(owner), state.Current); err != nil { return err @@ -64,13 +64,13 @@ func (r *runtimeLifecycle) wakeCompute(ctx context.Context, p sandbox.Checkpoint if err != nil { return err } - if err := r.store.ClearRuntimeWake(ctx, next, owner.ComputeActivityAt); err != nil { + if err := r.deployment.ClearWake(ctx, next, owner.ComputeActivityAt); err != nil { return err } return r.observeConnection(ctx, next) } -func (r *runtimeLifecycle) cleanupCompute(ctx context.Context, p sandbox.CheckpointProvider, owner store.RuntimeAllocation, state runtimeCompute) error { +func (r *runtimeLifecycle) cleanupCompute(ctx context.Context, p sandbox.CheckpointProvider, owner deployment.Allocation, state runtimeCompute) error { if err := r.lease.CheckOwnership(ctx); err != nil { return err } @@ -98,7 +98,7 @@ func (r *runtimeLifecycle) cleanupCompute(ctx context.Context, p sandbox.Checkpo return err } } - _, err := r.store.ReleaseRuntimeAllocation(ctx, owner) + _, err := r.deployment.ReleaseAllocation(ctx, owner) return err } @@ -108,8 +108,9 @@ func (w *Worker) waitRuntimeAwake(ctx context.Context, environment sessions.Envi if w.runtimes == nil { return nil } - owner, err := w.admission.GetRuntimeAllocation(ctx, environment.TenantID, environment.ID) - if errors.Is(err, sessions.ErrNotFound) { + key := deployment.AllocationKey{TenantID: environment.TenantID, EnvironmentID: environment.ID} + owner, err := w.dispatcher.DeploymentReader.EnvironmentAllocation(ctx, key) + if errors.Is(err, deployment.ErrNotFound) { return nil } if err != nil { @@ -121,13 +122,13 @@ func (w *Worker) waitRuntimeAwake(ctx context.Context, environment sessions.Envi if owner.ComputePhase == "disabled" { return nil } - if err := w.admission.TouchRuntimeActivity(ctx, environment.TenantID, environment.ID); err != nil { + if err := w.dispatcher.Deployment.TouchActivity(ctx, environment.TenantID, environment.ID); err != nil { return err } timer := time.NewTicker(100 * time.Millisecond) defer timer.Stop() for { - owner, err = w.admission.GetRuntimeAllocation(ctx, environment.TenantID, environment.ID) + owner, err = w.dispatcher.DeploymentReader.EnvironmentAllocation(ctx, key) if err != nil { return err } diff --git a/services/core/internal/execution/runtime_connections.go b/services/core/internal/execution/runtime_connections.go index bce757088..10c4bfd1a 100644 --- a/services/core/internal/execution/runtime_connections.go +++ b/services/core/internal/execution/runtime_connections.go @@ -6,9 +6,9 @@ import ( "github.com/google/uuid" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) // Access is serialized by the existing lifecycle gate. Durable generations fence @@ -20,7 +20,7 @@ type runtimeConnection struct { connected bool } -func (r *runtimeLifecycle) observeConnection(ctx context.Context, owner store.RuntimeAllocation) error { +func (r *runtimeLifecycle) observeConnection(ctx context.Context, owner deployment.Allocation) error { bound, err := r.sessions.GetSessionRuntimeDevice(ctx, owner.TenantID, owner.SessionID) if errors.Is(err, sessions.ErrNotFound) { return nil diff --git a/services/core/internal/execution/runtime_connections_test.go b/services/core/internal/execution/runtime_connections_test.go index fcbffafcb..c5b3d0624 100644 --- a/services/core/internal/execution/runtime_connections_test.go +++ b/services/core/internal/execution/runtime_connections_test.go @@ -3,12 +3,12 @@ package execution import ( "testing" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" ) func TestRuntimeCleanupDropsOnlyOwnedEnvironmentConnection(t *testing.T) { - first := store.RuntimeAllocation{ID: "allocation-one", EnvironmentID: "environment-one"} - second := store.RuntimeAllocation{ID: "allocation-two", EnvironmentID: "environment-two"} + first := deployment.Allocation{ID: "allocation-one", EnvironmentID: "environment-one"} + second := deployment.Allocation{ID: "allocation-two", EnvironmentID: "environment-two"} retained := &runtimeConnection{} r := &runtimeLifecycle{ connections: map[string]*runtimeConnection{first.EnvironmentID: {}, second.EnvironmentID: retained}, diff --git a/services/core/internal/execution/runtime_lifecycle.go b/services/core/internal/execution/runtime_lifecycle.go index f960edf04..54713e6d8 100644 --- a/services/core/internal/execution/runtime_lifecycle.go +++ b/services/core/internal/execution/runtime_lifecycle.go @@ -45,19 +45,24 @@ type runtimeLifecycle struct { store *store.Store sessions sessions.Reader sessionExecution *sessions.ExecutionOperations - lease Ownership - registry *runtimegateway.Registry - config RuntimeProvider - nodeID string - gate chan struct{} - ctx context.Context - stop context.CancelFunc - cancelMu sync.Mutex - reconcileCancel context.CancelFunc - cursor string - pendingCursor string - connections map[string]*runtimeConnection - wakeHints chan struct{} + // deployment runs the allocation mutations on the lease; deployments and + // reader run the pooled activity use case and reads. + deployment *deployment.ExecutionOperations + deployments *deployment.Service + reader deployment.Reader + lease Ownership + registry *runtimegateway.Registry + config RuntimeProvider + nodeID string + gate chan struct{} + ctx context.Context + stop context.CancelFunc + cancelMu sync.Mutex + reconcileCancel context.CancelFunc + cursor string + pendingCursor string + connections map[string]*runtimeConnection + wakeHints chan struct{} } func newRuntimeManager(owner Owner, deployments *deployment.Service, deploymentReader deployment.Reader, sessionReader sessions.Reader, registry *runtimegateway.Registry, config *RuntimeProvider) (*runtimeManager, error) { @@ -135,79 +140,80 @@ func (r *runtimeLifecycle) lock(ctx context.Context) error { // ProvisionEnvironment is an internal bootstrap operation for an already // authorized hosted Environment. It does not enable public hosted admission. -func (w *Worker) ProvisionEnvironment(ctx context.Context, tenant, environment, providerKey string) (store.RuntimeAllocation, error) { +func (w *Worker) ProvisionEnvironment(ctx context.Context, tenant, environment, providerKey string) (deployment.Allocation, error) { if w.runtimes == nil { - return store.RuntimeAllocation{}, ErrExecutionUnavailable + return deployment.Allocation{}, ErrExecutionUnavailable } ctx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() ctx, finish, err := w.runtimes.enter(ctx) if err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } defer finish() ready, err := w.runtimes.ensureDeployment(ctx) if err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } if !ready { - return store.RuntimeAllocation{}, ErrExecutionUnavailable + return deployment.Allocation{}, ErrExecutionUnavailable } - nodeID, err := w.runtimes.store.ResolveRuntimeLifecycleNode(ctx, tenant, environment) + nodeID, err := w.runtimes.deploymentService.LifecycleNode(ctx, tenant, environment) if err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } node, err := w.runtimes.node(nodeID) if err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } r := node.lifecycle if err := r.lock(ctx); err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } defer func() { <-r.gate }() return r.provision(ctx, tenant, environment, providerKey) } // provision runs under the lifecycle gate and uses the durable one-shot receipt. -func (r *runtimeLifecycle) provision(ctx context.Context, tenant, environment, providerKey string) (store.RuntimeAllocation, error) { +func (r *runtimeLifecycle) provision(ctx context.Context, tenant, environment, providerKey string) (deployment.Allocation, error) { provider := r.config.Provider if providerKey != r.config.InstallationID { - return store.RuntimeAllocation{}, sandbox.ErrInvalid + return deployment.Allocation{}, sandbox.ErrInvalid } environmentValue, err := r.sessions.GetEnvironment(ctx, tenant, environment) if err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } placement, err := parseEnvironmentPlacement(environmentValue.Configuration) if err != nil || placement.Type != "openai_hosted" { - return store.RuntimeAllocation{}, sandbox.ErrInvalid + return deployment.Allocation{}, sandbox.ErrInvalid } - if _, err := r.store.GetRuntimeAllocation(ctx, tenant, environment); errors.Is(err, sessions.ErrNotFound) { + key := deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment} + if _, err := r.reader.EnvironmentAllocation(ctx, key); errors.Is(err, deployment.ErrNotFound) { if r.config.AdmissionPaused && r.config.Generation == 0 { - return store.RuntimeAllocation{}, ErrExecutionUnavailable + return deployment.Allocation{}, ErrExecutionUnavailable } if err := r.computeFreshCapacity(ctx, providerKey); err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } if policy := r.config.Suspension; policy != nil && r.config.ProviderKind == "" { - count, err := r.store.CountRuntimeRetainedAllocations(ctx, providerKey) + count, err := r.reader.CountRetainedAllocations(ctx, providerKey) if err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } if count >= int64(policy.MaxRetained) { - return store.RuntimeAllocation{}, ErrExecutionUnavailable + return deployment.Allocation{}, ErrExecutionUnavailable } } } else if err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } secret := make([]byte, 32) if _, err := rand.Read(secret); err != nil { - return store.RuntimeAllocation{}, err + return deployment.Allocation{}, err } token := hex.EncodeToString(secret) - owner, err := r.store.ReserveRuntimeAllocation(ctx, tenant, environment, providerKey, runtimedevice.HashCredential(token)) + owner, err := r.deployment.ReserveAllocation(ctx, key, providerKey, runtimedevice.HashCredential(token)) if err != nil { return owner, err } @@ -226,7 +232,7 @@ func (r *runtimeLifecycle) provision(ctx context.Context, tenant, environment, p }) if info.Reference == runtimeReference(owner) && info.CreateSettled && info.State == "absent" { record, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) - released, releaseErr := r.store.ReleaseAbsentRuntimeCreation(record, owner) + released, releaseErr := r.deployment.ReleaseAbsentCreation(record, owner) cancel() if releaseErr != nil { return owner, releaseErr @@ -237,7 +243,7 @@ func (r *runtimeLifecycle) provision(ctx context.Context, tenant, environment, p return released, err } if info.Reference == runtimeReference(owner) && info.CreateSettled { - settled, settleErr := r.store.SettleRuntimeCreation(ctx, owner) + settled, settleErr := r.deployment.SettleCreation(ctx, owner) if settleErr != nil { return owner, settleErr } @@ -247,7 +253,7 @@ func (r *runtimeLifecycle) provision(ctx context.Context, tenant, environment, p // Explicit invalid/foreign bootstrap cannot become an authorized Runtime. // Other failures may hide a successful Create; retain recovery for those. if errors.Is(err, sandbox.ErrInvalid) || errors.Is(err, sandbox.ErrOwnership) { - if _, cleanupErr := r.store.RequestRuntimeCleanup(ctx, owner); cleanupErr != nil { + if _, cleanupErr := r.deployment.RequestCleanup(ctx, owner); cleanupErr != nil { return owner, cleanupErr } } @@ -257,7 +263,7 @@ func (r *runtimeLifecycle) provision(ctx context.Context, tenant, environment, p if info.Reference != runtimeReference(owner) || info.ProviderID == "" || info.State != "running" || !info.BootstrapComplete { return owner, sandbox.ErrOwnership } - return r.store.ObserveRuntimeRunning(ctx, owner) + return r.deployment.ObserveRunning(ctx, owner) } // ReconcileManagedRuntimes is also callable before serving admission. One scan @@ -276,7 +282,7 @@ func (r *runtimeLifecycle) reconcile(ctx context.Context) error { return err } defer finish() - rows, err := r.store.ListRuntimeAllocationsForNode(ctx, r.nodeID, r.cursor) + rows, err := r.reader.LifecycleAllocations(ctx, r.nodeID, r.cursor) if err != nil { return err } @@ -286,7 +292,7 @@ func (r *runtimeLifecycle) reconcile(ctx context.Context) error { if wrapped { // Service the next page now instead of spending a ticker interval on EOF. // Refill only once so an empty store still returns without spinning. - rows, err = r.store.ListRuntimeAllocationsForNode(ctx, r.nodeID, "") + rows, err = r.reader.LifecycleAllocations(ctx, r.nodeID, "") if err != nil { return err } @@ -310,13 +316,13 @@ func (r *runtimeLifecycle) reconcile(ctx context.Context) error { return r.provisionPending(ctx) } -func (r *runtimeLifecycle) observe(ctx context.Context, owner store.RuntimeAllocation) error { +func (r *runtimeLifecycle) observe(ctx context.Context, owner deployment.Allocation) error { if owner.ProviderKey != r.config.InstallationID || owner.NodeID != r.nodeID { return sandbox.ErrOwnership } if owner.SessionDeleted || owner.Expired || owner.State == "cleanup_pending" { var err error - owner, err = r.store.RequestRuntimeCleanup(ctx, owner) + owner, err = r.deployment.RequestCleanup(ctx, owner) if err != nil { return err } @@ -361,13 +367,13 @@ func (r *runtimeLifecycle) observe(ctx context.Context, owner store.RuntimeAlloc if (settled || running && info.BootstrapComplete) && !owner.CreateSettled { // Explicit settlement can accompany a configuration rejection. Missing // compute, a timeout or a successful Kill alone cannot settle Create. - owner, err = r.store.SettleRuntimeCreation(ctx, owner) + owner, err = r.deployment.SettleCreation(ctx, owner) if err != nil { return err } } if owner.SessionDeleted || owner.Expired || owner.State == "cleanup_pending" { - owner, err = r.store.RequestRuntimeCleanup(ctx, owner) + owner, err = r.deployment.RequestCleanup(ctx, owner) if err != nil { return err } @@ -380,7 +386,7 @@ func (r *runtimeLifecycle) observe(ctx context.Context, owner store.RuntimeAlloc if !owner.CreateSettled { return nil // Keep scanning unknown creation; absence is not a final receipt. } - _, err = r.store.ReleaseRuntimeAllocation(ctx, owner) + _, err = r.deployment.ReleaseAllocation(ctx, owner) return err } // A stopped or missing container does not authorize destroying retained @@ -394,7 +400,7 @@ func (r *runtimeLifecycle) observe(ctx context.Context, owner store.RuntimeAlloc } return nil } - owner, err = r.store.ObserveRuntimeRunning(ctx, owner) + owner, err = r.deployment.ObserveRunning(ctx, owner) if err != nil { return err } @@ -415,16 +421,16 @@ func (r *runtimeLifecycle) observe(ctx context.Context, owner store.RuntimeAlloc if renewed.Reference != runtimeReference(owner) || renewed.State != "running" || !renewed.BootstrapComplete { return sandbox.ErrOwnership } - _, err = r.store.KeepRuntimeAllocation(ctx, owner) + _, err = r.deployment.KeepAllocation(ctx, owner) return err } // Environment identity owns connectivity; preparation has an independent owner. -func (r *runtimeLifecycle) clearRuntimeState(owner store.RuntimeAllocation) { +func (r *runtimeLifecycle) clearRuntimeState(owner deployment.Allocation) { delete(r.connections, owner.EnvironmentID) } -func runtimeReference(owner store.RuntimeAllocation) sandbox.Reference { +func runtimeReference(owner deployment.Allocation) sandbox.Reference { return sandbox.Reference{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID, AllocationID: owner.ID} } diff --git a/services/core/internal/execution/runtime_manager.go b/services/core/internal/execution/runtime_manager.go index 589b1d95b..e2871bde3 100644 --- a/services/core/internal/execution/runtime_manager.go +++ b/services/core/internal/execution/runtime_manager.go @@ -89,7 +89,9 @@ func (m *runtimeManager) node(id string) (*runtimeNode, error) { if n == nil { ctx, stop := context.WithCancel(m.ctx) n = &runtimeNode{lifecycle: &runtimeLifecycle{ - store: m.store, sessions: m.sessions, sessionExecution: m.sessionExecution, lease: m.lease, registry: m.registry, config: m.config, nodeID: id, + store: m.store, sessions: m.sessions, sessionExecution: m.sessionExecution, + deployment: m.deployment, deployments: m.deploymentService, reader: m.deploymentReader, + 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), }} @@ -139,7 +141,7 @@ func (m *runtimeManager) syncNodes(ctx context.Context) ([]*runtimeNode, error) // A direct caller may add a newly registered node during the query. Only // entries present before this inventory snapshot can be retired by it. previous := m.snapshotNodes() - ids, err := m.store.ListRuntimeLifecycleNodes(ctx) + ids, err := m.deploymentReader.LifecycleNodes(ctx) if err != nil { return nil, err } diff --git a/services/core/internal/execution/runtime_observation.go b/services/core/internal/execution/runtime_observation.go index 82b5ccfd9..e4d93e3d0 100644 --- a/services/core/internal/execution/runtime_observation.go +++ b/services/core/internal/execution/runtime_observation.go @@ -4,11 +4,11 @@ import ( "context" "errors" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) -func (r *runtimeLifecycle) recordObservation(ctx context.Context, owner store.RuntimeAllocation, observed error) { +func (r *runtimeLifecycle) recordObservation(ctx context.Context, owner deployment.Allocation, observed error) { if owner.NodeID == "" { return } @@ -26,11 +26,11 @@ func (r *runtimeLifecycle) recordObservation(ctx context.Context, owner store.Ru diagnostic = "compute_unconfirmed" default: diagnostic = "provider_unavailable" - if online, err := r.store.RuntimeNodeAvailable(ctx, owner.NodeID); err == nil && !online { + if online, err := r.reader.NodeOnline(ctx, owner.NodeID); err == nil && !online { diagnostic = "node_unavailable" } } } // Reconciliation remains authoritative; diagnostics must not interrupt cleanup. - _ = r.store.RecordRuntimeObservation(ctx, owner, diagnostic) + _ = r.deployment.RecordObservation(ctx, owner, diagnostic) } diff --git a/services/core/internal/execution/runtime_pending.go b/services/core/internal/execution/runtime_pending.go index ebfb65a66..254b33877 100644 --- a/services/core/internal/execution/runtime_pending.go +++ b/services/core/internal/execution/runtime_pending.go @@ -13,7 +13,7 @@ func (r *runtimeLifecycle) provisionPending(ctx context.Context) error { if r.config.AdmissionPaused && r.config.Generation == 0 { return nil } - rows, err := r.store.ListUnallocatedHostedEnvironmentsForNode(ctx, r.nodeID, r.pendingCursor) + rows, err := r.reader.UnallocatedEnvironments(ctx, r.nodeID, r.pendingCursor) if err != nil { return err } diff --git a/services/core/internal/execution/runtime_wake_hint.go b/services/core/internal/execution/runtime_wake_hint.go index 074ce4c59..ed7c9f1d6 100644 --- a/services/core/internal/execution/runtime_wake_hint.go +++ b/services/core/internal/execution/runtime_wake_hint.go @@ -4,6 +4,7 @@ import ( "context" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -27,7 +28,7 @@ func (w *Worker) hintRuntimeWake(ctx context.Context, session sessions.Session) if err != nil || environment.Initialization != "complete" { return } - owner, err := w.admission.GetRuntimeAllocation(lookup, session.TenantID, environment.ID) + owner, err := w.dispatcher.DeploymentReader.EnvironmentAllocation(lookup, deployment.AllocationKey{TenantID: session.TenantID, EnvironmentID: environment.ID}) if err != nil || owner.ProviderKey != config.InstallationID || owner.State != "running" || !owner.CreateSettled || owner.SessionDeleted || owner.Expired { return diff --git a/services/core/internal/execution/sandbox_reset_test.go b/services/core/internal/execution/sandbox_reset_test.go index da24eec29..0f5212244 100644 --- a/services/core/internal/execution/sandbox_reset_test.go +++ b/services/core/internal/execution/sandbox_reset_test.go @@ -48,6 +48,7 @@ func resetManagerStoreDB(t *testing.T, configure func(*pgxpool.Config)) (*store. pool := pgtest.OpenIsolated(t, configure) cipher := testCredentialCipher(t) s := store.NewWithCredentialCipher(pool, cipher) + s.SetPlacement(fixtureRules(t)) owner, deployments, reader := testOwner(t, pool, cipher, s) return s, owner, deployments, reader, pool } diff --git a/services/core/internal/execution/worker_device.go b/services/core/internal/execution/worker_device.go index b17d83eb7..9e7cc9e16 100644 --- a/services/core/internal/execution/worker_device.go +++ b/services/core/internal/execution/worker_device.go @@ -7,6 +7,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -81,8 +82,8 @@ func (w *Worker) bindSessionDevice(ctx context.Context, session sessions.Session if _, err := parseEnvironmentPlacement(environment.Configuration); err != nil { return false, nil } - allocation, err := w.dispatcher.Store.GetRuntimeAllocation(ctx, session.TenantID, environment.ID) - if errors.Is(err, sessions.ErrNotFound) { + allocation, err := w.dispatcher.DeploymentReader.EnvironmentAllocation(ctx, deployment.AllocationKey{TenantID: session.TenantID, EnvironmentID: environment.ID}) + if errors.Is(err, deployment.ErrNotFound) { if snapshot.Environment.Type == "openai_hosted" { return false, nil } diff --git a/services/core/internal/persistence/postgres/deploymentpg/allocations.go b/services/core/internal/persistence/postgres/deploymentpg/allocations.go new file mode 100644 index 000000000..c12d7bca5 --- /dev/null +++ b/services/core/internal/persistence/postgres/deploymentpg/allocations.go @@ -0,0 +1,561 @@ +package deploymentpg + +import ( + "context" + "errors" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/placementpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// allocation converts a stored allocation with the facts of its Session. +func allocation(row sqlc.RuntimeAllocation, session, tenant pgtype.UUID, deleted pgtype.Timestamptz, expired bool) deployment.Allocation { + return deployment.Allocation{ + DeploymentGeneration: uint64(row.DeploymentGeneration.Int64), NodeID: uuidString(row.NodeID), ObservationError: row.ObservationError, + ComputePhase: row.ComputePhase, ComputeRevision: row.ComputeRevision, ComputeState: row.ComputeState, + ComputeActivityAt: row.ComputeActivityAt.Time, ComputeWakeRequested: row.ComputeWakeRequested, ComputeRetainedUntil: timestamp(row.ComputeRetainedUntil), + ID: uuidString(row.ID), EnvironmentID: uuidString(row.EnvironmentID), SessionID: uuidString(session), TenantID: uuidString(tenant), + DeviceID: uuidString(row.DeviceID), ProviderKey: uuidString(row.ProviderKey), State: row.State, CreateSettled: row.CreateSettled, + SessionDeleted: deleted.Valid, Expired: expired, CreatedAt: row.CreatedAt.Time, KeptAt: row.KeptAt.Time, + } +} + +// allocationKey returns the stored form of the key's identifiers. +func allocationKey(key deployment.AllocationKey) (pgtype.UUID, pgtype.UUID, error) { + tenant, err := parseID(key.TenantID) + if err != nil { + return tenant, pgtype.UUID{}, err + } + environment, err := parseID(key.EnvironmentID) + return tenant, environment, err +} + +// cursor returns the stored form of a page cursor: an empty one starts at +// the beginning. +func cursor(after string) (pgtype.UUID, error) { + if after == "" { + return pgtype.UUID{Valid: true}, nil + } + return parseID(after) +} + +// lifecycleNode returns the stored form of a lifecycle's node: an empty one +// names the lifecycle without a node. +func lifecycleNode(nodeID string) (pgtype.UUID, error) { + if nodeID == "" { + return pgtype.UUID{}, nil + } + return parseID(nodeID) +} + +func (e *Execution) WithReservation(ctx context.Context, key deployment.AllocationKey, apply func(sessions.LockedSession, deployment.ReservationTx) error) error { + tenant, environment, err := allocationKey(key) + if err != nil { + return err + } + return translate(e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + owned, err := q.GetEnvironment(ctx, sqlc.GetEnvironmentParams{TenantID: tenant, ID: environment}) + if errors.Is(err, pgx.ErrNoRows) { + // sessions.ErrNotFound keeps the public 404 that writeStoreError maps. + return sessions.ErrNotFound + } + if err != nil { + return err + } + session := owned.Environment.SessionID + locked, err := sessionpg.LockSession(ctx, q, tenant, session) + if err != nil { + return err + } + t := &reservationTx{SessionTx: sessionpg.BindSession(q, tenant, session), ctx: ctx, q: q, tenant: tenant, environment: environment, session: session} + if err := apply(locked, t); err != nil { + return err + } + return sessionpg.PruneChanges(ctx, q, session) + })) +} + +func (e *Execution) WithAllocation(ctx context.Context, key deployment.AllocationKey, apply func(deployment.AllocationTx) error) error { + return e.withAllocation(ctx, key, func(t *allocationTx) error { return apply(t) }) +} + +func (e *Execution) WithAllocationCleanup(ctx context.Context, key deployment.AllocationKey, apply func(deployment.AllocationCleanupTx) error) error { + return e.withAllocation(ctx, key, func(t *allocationTx) error { + return apply(&cleanupTx{allocationTx: t, SessionTx: sessionpg.BindSession(t.q, t.tenant, t.session)}) + }) +} + +// withAllocation locks the Session that owns the Environment's allocation, +// deleted or not, runs apply and prunes the Session's journal. +func (e *Execution) withAllocation(ctx context.Context, key deployment.AllocationKey, apply func(*allocationTx) error) error { + tenant, environment, err := allocationKey(key) + if err != nil { + return err + } + return translate(e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + row, err := q.GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: tenant, EnvironmentID: environment}) + if errors.Is(err, pgx.ErrNoRows) { + return deployment.ErrNotFound + } + if err != nil { + return err + } + if _, err := sessionpg.LockSession(ctx, q, row.TenantID, row.SessionID); err != nil { + return err + } + if err := apply(&allocationTx{ctx: ctx, q: q, tenant: tenant, environment: environment, session: row.SessionID}); err != nil { + return err + } + return sessionpg.PruneChanges(ctx, q, row.SessionID) + })) +} + +func (e *Execution) ClearWake(ctx context.Context, allocationID string, observed time.Time) error { + id, err := parseID(allocationID) + if err != nil { + return err + } + return e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + return sqlc.New(tx).ClearRuntimeWake(ctx, sqlc.ClearRuntimeWakeParams{ID: id, ComputeActivityAt: pgtype.Timestamptz{Time: observed, Valid: true}}) + }) +} + +func (s *Store) WithActivity(ctx context.Context, key deployment.AllocationKey, apply func(sessions.LockedSession, deployment.ActivityTx) error) error { + tenant, environment, err := allocationKey(key) + if err != nil { + return err + } + return translate(s.pool.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + owned, err := q.GetEnvironment(ctx, sqlc.GetEnvironmentParams{TenantID: tenant, ID: environment}) + if errors.Is(err, pgx.ErrNoRows) { + // sessions.ErrNotFound keeps the public 404 that writeStoreError maps. + return sessions.ErrNotFound + } + if err != nil { + return err + } + session := owned.Environment.SessionID + locked, err := sessionpg.LockSession(ctx, q, tenant, session) + if err != nil { + return err + } + touch := activityTx(func() error { + return q.TouchRuntimeActivity(ctx, sqlc.TouchRuntimeActivityParams{TenantID: tenant, EnvironmentID: environment}) + }) + if err := apply(locked, touch); err != nil { + return err + } + return sessionpg.PruneChanges(ctx, q, session) + })) +} + +// activityTx is one Session-locked activity change. +type activityTx func() error + +func (t activityTx) TouchActivity() error { return t() } + +// reservationTx is one Session-locked allocation reservation. +type reservationTx struct { + *sessionpg.SessionTx + ctx context.Context + q *sqlc.Queries + tenant, environment, session pgtype.UUID +} + +func (t *reservationTx) FindAllocation() (deployment.Allocation, bool, error) { + row, err := t.q.GetRuntimeAllocation(t.ctx, sqlc.GetRuntimeAllocationParams{TenantID: t.tenant, EnvironmentID: t.environment}) + if errors.Is(err, pgx.ErrNoRows) { + return deployment.Allocation{}, false, nil + } + if err != nil { + return deployment.Allocation{}, false, err + } + return allocation(row.RuntimeAllocation, t.session, t.tenant, row.DeletedAt, row.Expired), true, nil +} + +func (t *reservationTx) LockDeployment() (placement.Deployment, error) { + return placementpg.LockDeployment(t.ctx, t.q) +} + +func (t *reservationTx) LoadReserved() (placement.Reserved, error) { + return placementpg.LoadReserved(t.ctx, t.q, t.environment) +} + +func (t *reservationTx) InsertAllocation(a deployment.NewAllocation) (deployment.Allocation, error) { + id, err := parseID(a.ID) + if err != nil { + return deployment.Allocation{}, err + } + device, err := parseID(a.DeviceID) + if err != nil { + return deployment.Allocation{}, err + } + provider, err := parseID(a.ProviderKey) + if err != nil { + return deployment.Allocation{}, err + } + node, err := lifecycleNode(a.NodeID) + if err != nil { + return deployment.Allocation{}, err + } + row, err := t.q.CreateRuntimeAllocation(t.ctx, sqlc.CreateRuntimeAllocationParams{ + ID: id, EnvironmentID: t.environment, DeviceID: device, ProviderKey: provider, NodeID: node, + DeploymentGeneration: pgtype.Int8{Int64: int64(a.Generation), Valid: true}, + }) + if err != nil { + return deployment.Allocation{}, err + } + return allocation(row, t.session, t.tenant, pgtype.Timestamptz{}, false), nil +} + +// allocationTx is one allocation change in the transaction that locked the +// owning Session. +type allocationTx struct { + ctx context.Context + q *sqlc.Queries + tenant, environment, session pgtype.UUID +} + +func (t *allocationTx) LoadAllocation() (deployment.Allocation, error) { + row, err := t.q.GetRuntimeAllocation(t.ctx, sqlc.GetRuntimeAllocationParams{TenantID: t.tenant, EnvironmentID: t.environment}) + if errors.Is(err, pgx.ErrNoRows) { + return deployment.Allocation{}, deployment.ErrNotFound + } + if err != nil { + return deployment.Allocation{}, err + } + return allocation(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired), nil +} + +func (t *allocationTx) LoadSessionDevice() (deployment.SessionDevice, bool, error) { + row, err := t.q.GetSessionDevice(t.ctx, sqlc.GetSessionDeviceParams{TenantID: t.tenant, ID: t.session}) + if errors.Is(err, pgx.ErrNoRows) { + return deployment.SessionDevice{}, false, nil + } + if err != nil { + return deployment.SessionDevice{}, false, err + } + return deployment.SessionDevice{ID: uuidString(row.ID), EnvironmentID: uuidString(row.EnvironmentID)}, true, nil +} + +func (t *allocationTx) LoadActivity(current deployment.Allocation) (deployment.Activity, error) { + id, err := parseID(current.ID) + if err != nil { + return deployment.Activity{}, err + } + return loadActivity(t.ctx, t.q, id) +} + +func (t *allocationTx) LoadRestore(current deployment.Allocation) (placement.Restore, error) { + node, err := parseID(current.NodeID) + if err != nil { + return placement.Restore{}, err + } + return placementpg.LoadRestore(t.ctx, t.q, node, current.DeploymentGeneration) +} + +func (t *allocationTx) ObserveRunning(current deployment.Allocation) (deployment.Allocation, error) { + return t.change(current, t.q.ObserveRuntimeRunning) +} + +func (t *allocationTx) Keep(current deployment.Allocation) (deployment.Allocation, error) { + return t.change(current, t.q.KeepRuntimeAllocation) +} + +func (t *allocationTx) SettleCreation(current deployment.Allocation) (deployment.Allocation, error) { + return t.change(current, t.q.SettleRuntimeCreation) +} + +func (t *allocationTx) Release(current deployment.Allocation) (deployment.Allocation, error) { + released, err := t.change(current, t.q.ReleaseRuntimeAllocation) + if err != nil { + return deployment.Allocation{}, err + } + if err := placementpg.ReleasePlacement(t.ctx, t.q, t.environment); err != nil { + return deployment.Allocation{}, err + } + return released, nil +} + +func (t *allocationTx) SetCompute(current deployment.Allocation, change deployment.ComputeChange) (deployment.Allocation, error) { + until := pgtype.Timestamptz{} + if change.RetainedUntil != nil { + until = pgtype.Timestamptz{Time: *change.RetainedUntil, Valid: true} + } + return t.change(current, func(ctx context.Context, id pgtype.UUID) (sqlc.RuntimeAllocation, error) { + return t.q.SetRuntimeCompute(ctx, sqlc.SetRuntimeComputeParams{ID: id, Revision: current.ComputeRevision, Phase: change.Phase, State: change.State, RetainedUntil: until}) + }) +} + +func (t *allocationTx) RecordObservation(current deployment.Allocation, diagnostic string) error { + id, err := parseID(current.ID) + if err != nil { + return err + } + return t.q.SetRuntimeObservation(t.ctx, sqlc.SetRuntimeObservationParams{ID: id, ComputeRevision: current.ComputeRevision, State: current.State, ObservationError: diagnostic}) +} + +// change applies a guarded write to current and returns the stored result +// with current's Session facts. A guard that no longer holds is +// ErrAllocationConflict. +func (t *allocationTx) change(current deployment.Allocation, write func(context.Context, pgtype.UUID) (sqlc.RuntimeAllocation, error)) (deployment.Allocation, error) { + id, err := parseID(current.ID) + if err != nil { + return deployment.Allocation{}, err + } + row, err := write(t.ctx, id) + if errors.Is(err, pgx.ErrNoRows) { + return deployment.Allocation{}, deployment.ErrAllocationConflict + } + if err != nil { + return deployment.Allocation{}, err + } + changed := allocation(row, t.session, t.tenant, pgtype.Timestamptz{}, current.Expired) + changed.SessionDeleted = current.SessionDeleted + return changed, nil +} + +// cleanupTx is one allocation cleanup, with the owning Session's cancellation +// and Environment termination bound to the same transaction. +type cleanupTx struct { + *allocationTx + *sessionpg.SessionTx +} + +func (t *cleanupTx) RevokeDevice(current deployment.Allocation) error { + device, err := parseID(current.DeviceID) + if err != nil { + return err + } + _, err = t.q.RevokeRuntimeCleanupDevice(t.ctx, sqlc.RevokeRuntimeCleanupDeviceParams{TenantID: t.tenant, DeviceID: device}) + return err +} + +func (t *cleanupTx) RequestCleanup(current deployment.Allocation) (deployment.Allocation, error) { + return t.change(current, t.q.RequestRuntimeCleanup) +} + +func loadActivity(ctx context.Context, q *sqlc.Queries, id pgtype.UUID) (deployment.Activity, error) { + row, err := q.GetRuntimeActivity(ctx, id) + if errors.Is(err, pgx.ErrNoRows) { + return deployment.Activity{}, deployment.ErrNotFound + } + if err != nil { + return deployment.Activity{}, err + } + return deployment.Activity{LastActivity: row.LastActivity.Time, ObservedAt: row.ObservedAt.Time, Busy: row.Busy, WakeRequested: row.ComputeWakeRequested, HasCompletedTurn: row.HasCompletedTurn}, nil +} + +func (s *Store) EnvironmentAllocation(ctx context.Context, key deployment.AllocationKey) (deployment.Allocation, error) { + tenant, environment, err := allocationKey(key) + if err != nil { + return deployment.Allocation{}, err + } + row, err := s.pool.Queries().GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: tenant, EnvironmentID: environment}) + if errors.Is(err, pgx.ErrNoRows) { + return deployment.Allocation{}, deployment.ErrNotFound + } + if err != nil { + return deployment.Allocation{}, err + } + return allocation(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired), nil +} + +func (s *Store) CredentialAllocations(ctx context.Context, after string) ([]deployment.Allocation, error) { + id, err := cursor(after) + if err != nil { + return nil, err + } + rows, err := s.pool.Queries().ListRuntimeAllocations(ctx, id) + if err != nil { + return nil, err + } + result := make([]deployment.Allocation, 0, len(rows)) + for _, row := range rows { + result = append(result, allocation(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired)) + } + return result, nil +} + +func (s *Store) ObservationSessions(ctx context.Context, after string, limit int) (deployment.ObservationSessionPage, error) { + if limit < 1 || limit > 100 { + return deployment.ObservationSessionPage{}, deployment.ErrInvalidInput + } + id, err := cursor(after) + if err != nil { + return deployment.ObservationSessionPage{}, err + } + rows, err := s.pool.Queries().ListRuntimeObservationSessions(ctx, sqlc.ListRuntimeObservationSessionsParams{ID: id, Limit: int32(limit + 1)}) + if err != nil { + return deployment.ObservationSessionPage{}, err + } + page := deployment.ObservationSessionPage{Sessions: make([]deployment.ObservationSession, 0, min(limit, len(rows)))} + if len(rows) > limit { + page.NextCursor = uuidString(rows[limit-1].ID) + rows = rows[:limit] + } + for _, row := range rows { + page.Sessions = append(page.Sessions, deployment.ObservationSession{TenantID: uuidString(row.TenantID), SessionID: uuidString(row.ID)}) + } + return page, nil +} + +func (s *Store) NodeAllocations(ctx context.Context, nodeID string) ([]deployment.NodeAllocation, error) { + id, err := parseID(nodeID) + if err != nil { + return nil, err + } + q := s.pool.Queries() + if _, err := q.GetRuntimeNode(ctx, id); errors.Is(err, pgx.ErrNoRows) { + return nil, deployment.ErrNotFound + } else if err != nil { + return nil, err + } + rows, err := q.ListNodeRuntimeAllocations(ctx, id) + if err != nil { + return nil, err + } + result := make([]deployment.NodeAllocation, 0, len(rows)) + for _, a := range rows { + result = append(result, deployment.NodeAllocation{ + DeploymentGeneration: uint64(a.DeploymentGeneration.Int64), Diagnostic: a.ObservationError, ID: uuidString(a.ID), NodeID: uuidString(a.NodeID), + TenantID: uuidString(a.TenantID), SessionID: uuidString(a.SessionID), EnvironmentID: uuidString(a.EnvironmentID), State: a.State, + ComputePhase: a.ComputePhase, ComputePhaseChangedAt: timestamp(a.ComputePhaseChangedAt), Initialization: a.Initialization, CreatedAt: a.CreatedAt.Time, + }) + } + return result, nil +} + +func (s *Store) NodeOnline(ctx context.Context, nodeID string) (bool, error) { + id, err := parseID(nodeID) + if err != nil { + return false, err + } + nodes, err := s.pool.Queries().ListRuntimeNodes(ctx, pgtype.UUID{}) + if err != nil { + return false, err + } + for _, n := range nodes { + if n.ID == id { + return n.Online, nil + } + } + return false, nil +} + +// LifecycleNodes, LifecycleAllocations, UnallocatedEnvironments and +// LifecyclePlacement bound their reads by the execution deadline: the +// execution owner's lifecycles read through them before provider work. + +func (s *Store) LifecycleNodes(ctx context.Context) ([]string, error) { + ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) + defer cancel() + rows, err := s.pool.Queries().ListRuntimeLifecycleNodes(ctx) + if err != nil { + return nil, err + } + result := make([]string, 0, len(rows)) + for _, id := range rows { + result = append(result, uuidString(id)) + } + return result, nil +} + +func (s *Store) LifecycleAllocations(ctx context.Context, nodeID, after string) ([]deployment.Allocation, error) { + node, err := lifecycleNode(nodeID) + if err != nil { + return nil, err + } + id, err := cursor(after) + if err != nil { + return nil, err + } + ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) + defer cancel() + rows, err := s.pool.Queries().ListRuntimeAllocationsForNode(ctx, sqlc.ListRuntimeAllocationsForNodeParams{NodeID: node, AfterID: id}) + if err != nil { + return nil, err + } + result := make([]deployment.Allocation, 0, len(rows)) + for _, row := range rows { + result = append(result, allocation(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired)) + } + return result, nil +} + +func (s *Store) UnallocatedEnvironments(ctx context.Context, nodeID, after string) ([]deployment.UnallocatedEnvironment, error) { + node, err := lifecycleNode(nodeID) + if err != nil { + return nil, err + } + id, err := cursor(after) + if err != nil { + return nil, err + } + ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) + defer cancel() + rows, err := s.pool.Queries().ListUnallocatedHostedEnvironmentsForNode(ctx, sqlc.ListUnallocatedHostedEnvironmentsForNodeParams{NodeID: node, AfterID: id}) + if err != nil { + return nil, err + } + result := make([]deployment.UnallocatedEnvironment, 0, len(rows)) + for _, row := range rows { + result = append(result, deployment.UnallocatedEnvironment{ID: uuidString(row.ID), TenantID: uuidString(row.TenantID)}) + } + return result, nil +} + +func (s *Store) LifecyclePlacement(ctx context.Context, key deployment.AllocationKey) (deployment.LifecyclePlacement, error) { + tenant, environment, err := allocationKey(key) + if err != nil { + return deployment.LifecyclePlacement{}, err + } + ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) + defer cancel() + row, err := s.pool.Queries().GetRuntimeLifecyclePlacement(ctx, sqlc.GetRuntimeLifecyclePlacementParams{TenantID: tenant, ID: environment}) + if errors.Is(err, pgx.ErrNoRows) { + // sessions.ErrNotFound keeps the public 404 that writeStoreError maps. + return deployment.LifecyclePlacement{}, sessions.ErrNotFound + } + if err != nil { + return deployment.LifecyclePlacement{}, err + } + return deployment.LifecyclePlacement{Provider: row.ProviderKind, Mode: row.Mode, AllocationID: uuidString(row.AllocationID), AllocationNodeID: uuidString(row.AllocationNodeID), + PlacementNodeID: uuidString(row.PlacementNodeID), PlacementReleased: row.ReleasedAt.Valid}, nil +} + +func (s *Store) Activity(ctx context.Context, allocationID string) (deployment.Activity, error) { + id, err := parseID(allocationID) + if err != nil { + return deployment.Activity{}, err + } + return loadActivity(ctx, s.pool.Queries(), id) +} + +func (s *Store) CountComputeReservations(ctx context.Context, installationID string) (int64, error) { + id, err := parseID(installationID) + if err != nil { + return 0, err + } + return s.pool.Queries().CountRuntimeComputeReservations(ctx, id) +} + +func (s *Store) CountRetainedAllocations(ctx context.Context, installationID string) (int64, error) { + id, err := parseID(installationID) + if err != nil { + return 0, err + } + return s.pool.Queries().CountRuntimeRetainedAllocations(ctx, id) +} diff --git a/services/core/internal/persistence/postgres/deploymentpg/allocations_test.go b/services/core/internal/persistence/postgres/deploymentpg/allocations_test.go new file mode 100644 index 000000000..0fed81384 --- /dev/null +++ b/services/core/internal/persistence/postgres/deploymentpg/allocations_test.go @@ -0,0 +1,203 @@ +package deploymentpg_test + +import ( + "crypto/sha256" + "encoding/hex" + "errors" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" +) + +// hostedEnvironment stores a fresh tenant's hosted Session with its pending +// Environment and returns the Environment's key. +func hostedEnvironment(t *testing.T, pool *pgxpool.Pool) deployment.AllocationKey { + t.Helper() + tenant, session, environment := uuid.New(), uuid.New(), uuid.New() + if _, err := pool.Exec(t.Context(), `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash, configuration) + VALUES ($1, $2, 'codex', 'key', 'hash', '{"environment":{"type":"openai_hosted"}}')`, session, tenant); err != nil { + t.Fatal(err) + } + if _, err := pool.Exec(t.Context(), `INSERT INTO environments(id, session_id) VALUES ($1, $2)`, environment, session); err != nil { + t.Fatal(err) + } + return deployment.AllocationKey{TenantID: tenant.String(), EnvironmentID: environment.String()} +} + +// credentialHash is the digest of a fresh device credential. +func credentialHash() string { + digest := sha256.Sum256([]byte(uuid.NewString())) + return hex.EncodeToString(digest[:]) +} + +// allocationRows reports the Environment's allocation state, device +// revocation and status, and the Session's journal length. +func allocationRows(t *testing.T, f fixture, key deployment.AllocationKey) (allocations int, state string, revoked bool, status string, changes int) { + t.Helper() + err := f.pool.QueryRow(t.Context(), `SELECT + (SELECT count(*) FROM runtime_allocations WHERE environment_id = e.id), + COALESCE((SELECT state FROM runtime_allocations WHERE environment_id = e.id), ''), + COALESCE((SELECT revoked_at IS NOT NULL FROM devices WHERE environment_id = e.id), false), + e.status, + (SELECT count(*) FROM session_events WHERE session_id = e.session_id) + FROM environments e WHERE e.id = $1`, key.EnvironmentID).Scan(&allocations, &state, &revoked, &status, &changes) + if err != nil { + t.Fatal(err) + } + return +} + +// Concurrent reservations of one Environment serialize on its Session: one +// commits the allocation and its device, the others replay it, and another +// installation conflicts. +func TestConcurrentReservationsCommitOneAllocation(t *testing.T) { + f := newFixture(t) + changes, _ := f.execution(t) + installation, _ := f.initialize(t, changes, setupE2BSelection()) + key := hostedEnvironment(t, f.pool) + results := make([]deployment.Allocation, 8) + errs := make([]error, len(results)) + var wg sync.WaitGroup + for i := range results { + wg.Go(func() { + results[i], errs[i] = changes.ReserveAllocation(t.Context(), key, installation, credentialHash()) + }) + } + wg.Wait() + fresh := 0 + for i, result := range results { + if errs[i] != nil || result.ID != results[0].ID || result.DeviceID != results[0].DeviceID { + t.Fatal("reservations disagree", result, errs[i]) + } + if !result.Replayed { + fresh++ + } + } + var devices int + if err := f.pool.QueryRow(t.Context(), "SELECT count(*) FROM devices WHERE environment_id = $1", key.EnvironmentID).Scan(&devices); err != nil { + t.Fatal(err) + } + if allocations, state, _, _, _ := allocationRows(t, f, key); fresh != 1 || allocations != 1 || devices != 1 || state != "creating" { + t.Fatal("reservations committed more than one allocation", fresh, allocations, devices, state) + } + if _, err := changes.ReserveAllocation(t.Context(), key, uuid.NewString(), credentialHash()); !errors.Is(err, deployment.ErrAllocationConflict) { + t.Fatal("another installation replayed the allocation", err) + } +} + +// Every allocation write needs the execution lease; after the lease is lost +// none of them commits. +func TestAllocationWritesNeedTheLease(t *testing.T) { + f := newFixture(t) + closed := setupClosedExecution(t, f) + changes, _ := f.execution(t) + installation, _ := f.initialize(t, changes, setupE2BSelection()) + key := hostedEnvironment(t, f.pool) + owner, err := changes.ReserveAllocation(t.Context(), key, installation, credentialHash()) + if err != nil { + t.Fatal(err) + } + unallocated := hostedEnvironment(t, f.pool) + if _, err := closed.ReserveAllocation(t.Context(), unallocated, installation, credentialHash()); !errors.Is(err, pgunit.ErrLeaseClosed) { + t.Fatal("reserved without the lease", err) + } + if allocations, _, _, _, _ := allocationRows(t, f, unallocated); allocations != 0 { + t.Fatal("a reservation without the lease committed") + } + writes := map[string]func() error{ + "ObserveRunning": func() error { _, err := closed.ObserveRunning(t.Context(), owner); return err }, + "SettleCreation": func() error { _, err := closed.SettleCreation(t.Context(), owner); return err }, + "RequestCleanup": func() error { _, err := closed.RequestCleanup(t.Context(), owner); return err }, + "SetCompute": func() error { + _, err := closed.SetCompute(t.Context(), owner, "running", []byte(`{}`), nil, 0) + return err + }, + "ClearWake": func() error { return closed.ClearWake(t.Context(), owner, time.Now()) }, + } + for name, write := range writes { + if err := write(); !errors.Is(err, pgunit.ErrLeaseClosed) { + t.Fatal(name, "wrote without the lease", err) + } + } + if current, err := f.adapter.EnvironmentAllocation(t.Context(), key); err != nil || current.State != "creating" || current.CreateSettled || current.ComputePhase != owner.ComputePhase { + t.Fatal("a write without the lease committed", current, err) + } +} + +// Cleanup revokes the device, settles the Session and requests cleanup in +// one transaction: a failed settlement rolls the revocation back. +func TestFailedCleanupSettlementRollsBackRevocation(t *testing.T) { + f := newFixture(t) + changes, _ := f.execution(t) + installation, _ := f.initialize(t, changes, setupE2BSelection()) + key := hostedEnvironment(t, f.pool) + owner, err := changes.ReserveAllocation(t.Context(), key, installation, credentialHash()) + if err != nil { + t.Fatal(err) + } + if owner, err = changes.ObserveRunning(t.Context(), owner); err != nil { + t.Fatal(err) + } + _, _, _, _, before := allocationRows(t, f, key) + if _, err := f.pool.Exec(t.Context(), `CREATE FUNCTION fail_environment_update() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'injected settlement failure'; END $$; + CREATE TRIGGER fail_environment_update BEFORE UPDATE ON environments FOR EACH ROW EXECUTE FUNCTION fail_environment_update()`); err != nil { + t.Fatal(err) + } + if _, err := changes.RequestCleanup(t.Context(), owner); err == nil { + t.Fatal("cleanup committed without the Session settlement") + } + if _, state, revoked, status, changed := allocationRows(t, f, key); state != "running" || revoked || status != "pending" || changed != before { + t.Fatal("failed cleanup left a partial commit", state, revoked, status, changed) + } + if _, err := f.pool.Exec(t.Context(), "DROP TRIGGER fail_environment_update ON environments"); err != nil { + t.Fatal(err) + } + pending, err := changes.RequestCleanup(t.Context(), owner) + if err != nil || pending.State != "cleanup_pending" { + t.Fatal(pending, err) + } + if _, state, revoked, status, changed := allocationRows(t, f, key); state != "cleanup_pending" || !revoked || status != "failed" || changed <= before { + t.Fatal("cleanup did not commit together", state, revoked, status, changed) + } +} + +// An allocation write commits only with the Session journal prune: a failed +// prune rolls back a change and a reservation with its device. +func TestFailedPruneRollsBackTheAllocationWrite(t *testing.T) { + f := newFixture(t) + changes, _ := f.execution(t) + installation, _ := f.initialize(t, changes, setupE2BSelection()) + key := hostedEnvironment(t, f.pool) + owner, err := changes.ReserveAllocation(t.Context(), key, installation, credentialHash()) + if err != nil { + t.Fatal(err) + } + if _, err := f.pool.Exec(t.Context(), `CREATE FUNCTION fail_prune() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'injected prune failure'; END $$; + CREATE TRIGGER fail_prune BEFORE DELETE ON session_events FOR EACH STATEMENT EXECUTE FUNCTION fail_prune()`); err != nil { + t.Fatal(err) + } + if _, err := changes.ObserveRunning(t.Context(), owner); err == nil || !strings.Contains(err.Error(), "injected prune failure") { + t.Fatal("a change committed without the prune", err) + } + if _, state, _, _, _ := allocationRows(t, f, key); state != "creating" { + t.Fatal("a failed prune kept the change", state) + } + unallocated := hostedEnvironment(t, f.pool) + if _, err := changes.ReserveAllocation(t.Context(), unallocated, installation, credentialHash()); err == nil || !strings.Contains(err.Error(), "injected prune failure") { + t.Fatal("a reservation committed without the prune", err) + } + var devices int + if err := f.pool.QueryRow(t.Context(), "SELECT count(*) FROM devices WHERE environment_id = $1", unallocated.EnvironmentID).Scan(&devices); err != nil { + t.Fatal(err) + } + if allocations, _, _, _, _ := allocationRows(t, f, unallocated); allocations != 0 || devices != 0 { + t.Fatal("a failed prune kept the reservation", allocations, devices) + } +} diff --git a/services/core/internal/persistence/postgres/deploymentpg/fixture_test.go b/services/core/internal/persistence/postgres/deploymentpg/fixture_test.go index 035a607be..45c48d97a 100644 --- a/services/core/internal/persistence/postgres/deploymentpg/fixture_test.go +++ b/services/core/internal/persistence/postgres/deploymentpg/fixture_test.go @@ -12,6 +12,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/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" @@ -47,11 +48,23 @@ func newFixture(t *testing.T) fixture { func (f fixture) withPublicURL(t *testing.T, publicURL string) (*deploymentpg.Store, *deployment.Service) { t.Helper() adapter := deploymentpg.New(pgunit.NewPool(f.pool), f.cipher) - service, err := deployment.NewService(adapter, adapter, providers.Builtin(), publicURL) + return adapter, newService(t, adapter, publicURL) +} + +// newService builds the deployment service on adapter, with the placement +// rules for this public URL, as cmd/server does. +func newService(t *testing.T, adapter *deploymentpg.Store, publicURL string) *deployment.Service { + t.Helper() + registry := providers.Builtin() + rules, err := placement.NewRules(registry, publicURL) + if err != nil { + t.Fatal(err) + } + service, err := deployment.NewService(adapter, adapter, registry, rules) if err != nil { t.Fatal(err) } - return adapter, service + return service } // execution acquires the execution lease and builds the deployment execution diff --git a/services/core/internal/persistence/postgres/deploymentpg/presence_test.go b/services/core/internal/persistence/postgres/deploymentpg/presence_test.go index 62eb23d68..7787174dd 100644 --- a/services/core/internal/persistence/postgres/deploymentpg/presence_test.go +++ b/services/core/internal/persistence/postgres/deploymentpg/presence_test.go @@ -17,7 +17,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" ) func presenceContext(t *testing.T) context.Context { @@ -165,10 +164,7 @@ func TestNodePresenceCanceledBeforeCommit(t *testing.T) { } defer pool.Close() adapter := deploymentpg.New(pgunit.NewPool(pool), f.cipher) - service, err := deployment.NewService(adapter, adapter, providers.Builtin(), fixturePublicURL) - if err != nil { - t.Fatal(err) - } + service := newService(t, adapter, fixturePublicURL) if err := service.ConnectNode(ctx, node, uuid.NewString(), epoch); !errors.Is(err, context.Canceled) { t.Fatal("canceled UPDATE published presence", err) } diff --git a/services/core/internal/persistence/postgres/deploymentpg/seal_test.go b/services/core/internal/persistence/postgres/deploymentpg/seal_test.go index da1186920..f52e389ca 100644 --- a/services/core/internal/persistence/postgres/deploymentpg/seal_test.go +++ b/services/core/internal/persistence/postgres/deploymentpg/seal_test.go @@ -8,7 +8,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" ) // A missing key is credentialcrypto.ErrUnavailable, and a credential sealed @@ -19,10 +18,7 @@ func TestStoredCredentialNeedsTheKeyAndItsBinding(t *testing.T) { changes, _ := f.execution(t) id, view := f.initialize(t, changes, setupE2BSelection()) keyless := deploymentpg.New(pgunit.NewPool(f.pool), nil) - service, err := deployment.NewService(keyless, keyless, providers.Builtin(), fixturePublicURL) - if err != nil { - t.Fatal(err) - } + service := newService(t, keyless, fixturePublicURL) if _, err := service.Setup(t.Context()); !errors.Is(err, credentialcrypto.ErrUnavailable) { t.Fatal("a keyless Store opened the credential", err) } diff --git a/services/core/internal/persistence/postgres/placementpg/placementpg.go b/services/core/internal/persistence/postgres/placementpg/placementpg.go new file mode 100644 index 000000000..f67cbf294 --- /dev/null +++ b/services/core/internal/persistence/postgres/placementpg/placementpg.go @@ -0,0 +1,109 @@ +// Package placementpg loads the placement facts of the sandbox deployment and +// applies the placements package placement decides, on the caller's +// transaction-bound queries. It has no Store, pool or transaction runner: +// deploymentpg, sessionpg and Session creation call it inside their own +// transactions, which supply the authority. +package placementpg + +import ( + "context" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" +) + +// LockDeployment locks the deployment and returns it as placement reads it. +// Admission, placement, restore and node removal serialize on this lock. +func LockDeployment(ctx context.Context, q *sqlc.Queries) (placement.Deployment, error) { + d, err := q.LockRuntimeDeployment(ctx) + if err != nil { + return placement.Deployment{}, err + } + return placement.Deployment{ + InstallationID: uuidString(d.InstallationID), Provider: d.ProviderKind, Mode: d.Mode, + Generation: uint64(d.Generation), WebManaged: d.WebManaged, AdmissionPaused: d.AdmissionPaused, + Resetting: d.ResetClear.Valid, Specification: d.Specification, + }, nil +} + +// LoadNodes returns the installation's nodes with their presence and usage. +func LoadNodes(ctx context.Context, q *sqlc.Queries) ([]placement.Node, error) { + rows, err := q.ListRuntimeNodes(ctx, pgtype.UUID{}) + if err != nil { + return nil, err + } + nodes := make([]placement.Node, 0, len(rows)) + for _, row := range rows { + nodes = append(nodes, node(row)) + } + return nodes, nil +} + +// ReservePlacement reserves the placement for the Session's Environment. +func ReservePlacement(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, p placement.Placement) error { + id, err := pgunit.ParseID(p.NodeID) + if err != nil { + return err + } + return q.CreateSessionRuntimePlacement(ctx, sqlc.CreateSessionRuntimePlacementParams{SessionID: session, NodeID: id, Generation: int64(p.Generation)}) +} + +// LoadReserved returns the node placement reserved for the Environment. +func LoadReserved(ctx context.Context, q *sqlc.Queries, environment pgtype.UUID) (placement.Reserved, error) { + row, err := q.GetRuntimePlacement(ctx, environment) + if err != nil { + return placement.Reserved{}, err + } + return placement.Reserved{NodeID: uuidString(row.NodeID), Generation: uint64(row.DeploymentGeneration.Int64), Released: row.ReleasedAt.Valid, Available: row.Available}, nil +} + +// ReleasePlacement releases the node placement of the Environment, if any. +func ReleasePlacement(ctx context.Context, q *sqlc.Queries, environment pgtype.UUID) error { + return q.ReleaseRuntimePlacement(ctx, environment) +} + +// LoadRestore locks the deployment and returns what restoring suspended +// compute of the generation on the node reads. +func LoadRestore(ctx context.Context, q *sqlc.Queries, nodeID pgtype.UUID, generation uint64) (placement.Restore, error) { + if _, err := q.LockRuntimeDeployment(ctx); err != nil { + return placement.Restore{}, err + } + restore := placement.Restore{Generation: generation} + rows, err := q.ListRuntimeNodes(ctx, nodeID) + if err != nil || len(rows) == 0 { + return restore, err + } + n := node(rows[0]) + restore.Node = &n + restore.GenerationReady, err = q.NodeGenerationReady(ctx, sqlc.NodeGenerationReadyParams{NodeID: nodeID, Generation: int64(generation)}) + return restore, err +} + +// ComputeBlocksAdmission reports whether the Session's managed compute is in +// its suspension cycle: quiescing, suspending, suspended, restoring or waking. +func ComputeBlocksAdmission(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) (bool, error) { + return q.RuntimeComputeBlocksAdmission(ctx, session) +} + +func node(row sqlc.ListRuntimeNodesRow) placement.Node { + n := placement.Node{ + ID: uuidString(row.ID), Online: row.Online, ServingReady: row.ServingReady, TargetState: row.TargetState, + Active: row.Active, Retained: row.Retained, MaxActive: int(row.MaxActive), MaxRetained: int(row.MaxRetained), CoreURL: row.CoreUrl, + } + if row.ReadyGeneration.Valid { + ready := uint64(row.ReadyGeneration.Int64) + n.ReadyGeneration = &ready + } + return n +} + +func uuidString(id pgtype.UUID) string { + if !id.Valid { + return "" + } + return uuid.UUID(id.Bytes).String() +} diff --git a/services/core/internal/persistence/postgres/sessionpg/binding.go b/services/core/internal/persistence/postgres/sessionpg/binding.go index 085e99613..04fadf5bd 100644 --- a/services/core/internal/persistence/postgres/sessionpg/binding.go +++ b/services/core/internal/persistence/postgres/sessionpg/binding.go @@ -10,6 +10,7 @@ import ( "github.com/jackc/pgx/v5/pgtype" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/placementpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -187,7 +188,7 @@ func (t *SessionTx) InsertEnvironmentDevice(ctx context.Context, device sessions } func (t *SessionTx) LoadComputeSuspension(ctx context.Context) (bool, error) { - return t.q.RuntimeComputeBlocksAdmission(ctx, t.session) + return placementpg.ComputeBlocksAdmission(ctx, t.q, t.session) } func (t *SessionTx) LoadPendingFileWrite(ctx context.Context) (bool, error) { diff --git a/services/core/internal/runtimeobs/storeresolver/resolver.go b/services/core/internal/runtimeobs/storeresolver/resolver.go index 817378085..8f7e8e8dd 100644 --- a/services/core/internal/runtimeobs/storeresolver/resolver.go +++ b/services/core/internal/runtimeobs/storeresolver/resolver.go @@ -8,25 +8,32 @@ import ( "fmt" "io" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) type sessionStore interface { GetSession(context.Context, string, string) (sessions.Session, error) MeasuredSessionUsage(context.Context, string, string) (json.RawMessage, error) - GetRuntimeAllocation(context.Context, string, string) (store.RuntimeAllocation, error) - ListRuntimeObservationSessions(context.Context, string, int) (store.RuntimeObservationSessionPage, error) } -type Resolver struct{ store sessionStore } +// allocationReader is the part of deployment.Reader the resolver reads. +type allocationReader interface { + EnvironmentAllocation(context.Context, deployment.AllocationKey) (deployment.Allocation, error) + ObservationSessions(context.Context, string, int) (deployment.ObservationSessionPage, error) +} + +type Resolver struct { + store sessionStore + allocations allocationReader +} -func NewResolver(s sessionStore) (*Resolver, error) { - if s == nil { +func NewResolver(s sessionStore, allocations allocationReader) (*Resolver, error) { + if s == nil || allocations == nil { return nil, errors.New("Runtime observation store is required") } - return &Resolver{store: s}, nil + return &Resolver{store: s, allocations: allocations}, nil } func (r *Resolver) Resolve(ctx context.Context, tenantID, sessionID string) (runtimeobs.Target, error) { @@ -77,8 +84,8 @@ func (r *Resolver) Resolve(ctx context.Context, tenantID, sessionID string) (run return runtimeobs.Target{}, errors.New("managed Environment does not match resolved ownership") } target.EnvironmentID = session.Environment.ID - allocation, err := r.store.GetRuntimeAllocation(ctx, tenantID, target.EnvironmentID) - if errors.Is(err, sessions.ErrNotFound) { + allocation, err := r.allocations.EnvironmentAllocation(ctx, deployment.AllocationKey{TenantID: tenantID, EnvironmentID: target.EnvironmentID}) + if errors.Is(err, deployment.ErrNotFound) { return target, runtimeobs.ErrUnavailable } if err != nil { @@ -147,7 +154,7 @@ func decodeTokenUsage(raw json.RawMessage) (*runtimeobs.TokenUsage, error) { } func (r *Resolver) ListRuntimeObservationSessions(ctx context.Context, after string, limit int) (runtimeobs.SessionPage, error) { - page, err := r.store.ListRuntimeObservationSessions(ctx, after, limit) + page, err := r.allocations.ObservationSessions(ctx, after, limit) if err != nil { return runtimeobs.SessionPage{}, err } diff --git a/services/core/internal/runtimeobs/storeresolver/resolver_test.go b/services/core/internal/runtimeobs/storeresolver/resolver_test.go index e9c028d9e..593d6283a 100644 --- a/services/core/internal/runtimeobs/storeresolver/resolver_test.go +++ b/services/core/internal/runtimeobs/storeresolver/resolver_test.go @@ -6,17 +6,17 @@ import ( "errors" "testing" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) type resolverStore struct { session sessions.Session measured json.RawMessage - allocation store.RuntimeAllocation + allocation deployment.Allocation allocationErr error - page store.RuntimeObservationSessionPage + page deployment.ObservationSessionPage } func (s resolverStore) GetSession(context.Context, string, string) (sessions.Session, error) { @@ -30,17 +30,20 @@ func (s resolverStore) MeasuredSessionUsage(_ context.Context, tenant, session s return s.measured, nil } -func (s resolverStore) GetRuntimeAllocation(context.Context, string, string) (store.RuntimeAllocation, error) { +func (s resolverStore) EnvironmentAllocation(context.Context, deployment.AllocationKey) (deployment.Allocation, error) { return s.allocation, s.allocationErr } -func (s resolverStore) ListRuntimeObservationSessions(context.Context, string, int) (store.RuntimeObservationSessionPage, error) { +func (s resolverStore) ObservationSessions(context.Context, string, int) (deployment.ObservationSessionPage, error) { return s.page, nil } +// newTestResolver serves the Session and the deployment reads from one fake. +func newTestResolver(s resolverStore) (*Resolver, error) { return NewResolver(s, s) } + func TestResolverListsOnlyProviderNeutralSessionIdentity(t *testing.T) { - r, err := NewResolver(resolverStore{page: store.RuntimeObservationSessionPage{ - Sessions: []store.RuntimeObservationSession{{TenantID: "tenant", SessionID: "session"}}, + r, err := newTestResolver(resolverStore{page: deployment.ObservationSessionPage{ + Sessions: []deployment.ObservationSession{{TenantID: "tenant", SessionID: "session"}}, NextCursor: "session", }}) if err != nil { @@ -53,10 +56,10 @@ func TestResolverListsOnlyProviderNeutralSessionIdentity(t *testing.T) { } func TestResolverBindsManagedSessionEnvironmentAndAllocation(t *testing.T) { - r, err := NewResolver(resolverStore{ + r, err := newTestResolver(resolverStore{ session: sessions.Session{ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"openai_hosted"}}`), Environment: &sessions.Environment{ID: "environment", TenantID: "tenant", SessionID: "session"}}, measured: []byte(`{"input_tokens":120,"input_tokens_details":{"cached_tokens":20},"output_tokens":30,"output_tokens_details":{"reasoning_tokens":10},"total_tokens":150}`), - allocation: store.RuntimeAllocation{ + allocation: deployment.Allocation{ ID: "allocation", TenantID: "tenant", SessionID: "session", EnvironmentID: "environment", ProviderKey: "provider", DeviceID: "device", ComputePhase: "running", ComputeState: []byte(`{"current":{"name":"sandbox"}}`), }, @@ -90,7 +93,7 @@ func TestResolverRejectsInvalidCanonicalSessionUsage(t *testing.T) { `{"input_tokens":0,"input_tokens_details":{"cached_tokens":-1},"output_tokens":0,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":0}`, `{"input_tokens":0,"input_tokens_details":{"cached_tokens":0},"output_tokens":0,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":0,"unknown":0}`, } { - resolver, err := NewResolver(resolverStore{session: sessions.Session{ + resolver, err := newTestResolver(resolverStore{session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"none"}}`), }, measured: []byte(usage)}) if err != nil { @@ -103,7 +106,7 @@ func TestResolverRejectsInvalidCanonicalSessionUsage(t *testing.T) { } func TestResolverKeepsNullCanonicalSessionUsageAbsent(t *testing.T) { - resolver, err := NewResolver(resolverStore{session: sessions.Session{ + resolver, err := newTestResolver(resolverStore{session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"none"}}`), }, measured: []byte(" \n null \t")}) if err != nil { @@ -126,7 +129,7 @@ func TestResolverUsesMeasuredRatherThanPublicSessionUsage(t *testing.T) { {"null", measured, &runtimeobs.TokenUsage{InputTokens: 7, OutputTokens: 3}}, {measured, "null", nil}, } { - resolver, err := NewResolver(resolverStore{session: sessions.Session{ + resolver, err := newTestResolver(resolverStore{session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"none"}}`), Usage: []byte(test.public), }, measured: []byte(test.measured)}) if err != nil { @@ -147,7 +150,7 @@ func TestResolverKeepsUnsupportedModesDistinct(t *testing.T) { {mode: "none"}, {mode: "self_hosted", environment: &sessions.Environment{ID: "environment", TenantID: "tenant", SessionID: "session"}}, } { - r, err := NewResolver(resolverStore{session: sessions.Session{ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"` + tc.mode + `"}}`), Environment: tc.environment}}) + r, err := newTestResolver(resolverStore{session: sessions.Session{ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"` + tc.mode + `"}}`), Environment: tc.environment}}) if err != nil { t.Fatal(err) } @@ -159,9 +162,9 @@ func TestResolverKeepsUnsupportedModesDistinct(t *testing.T) { } func TestResolverReportsManagedAllocationAsUnavailable(t *testing.T) { - r, err := NewResolver(resolverStore{ + r, err := newTestResolver(resolverStore{ session: sessions.Session{ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"openai_hosted"}}`), Environment: &sessions.Environment{ID: "environment", TenantID: "tenant", SessionID: "session"}}, - allocationErr: sessions.ErrNotFound, + allocationErr: deployment.ErrNotFound, }) if err != nil { t.Fatal(err) @@ -178,7 +181,7 @@ func TestResolverRejectsMismatchedEnvironmentOwnership(t *testing.T) { {ID: "environment", TenantID: "other", SessionID: "session"}, {ID: "environment", TenantID: "tenant", SessionID: "other"}, } { - resolver, err := NewResolver(resolverStore{session: sessions.Session{ + resolver, err := newTestResolver(resolverStore{session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"` + mode + `"}}`), Environment: &environment, }}) @@ -193,18 +196,18 @@ func TestResolverRejectsMismatchedEnvironmentOwnership(t *testing.T) { } func TestResolverRejectsMismatchedAllocationOwnership(t *testing.T) { - base := store.RuntimeAllocation{ + base := deployment.Allocation{ ID: "allocation", TenantID: "tenant", SessionID: "session", EnvironmentID: "environment", ProviderKey: "provider", DeviceID: "device", } - for _, mutate := range []func(*store.RuntimeAllocation){ - func(value *store.RuntimeAllocation) { value.TenantID = "other" }, - func(value *store.RuntimeAllocation) { value.SessionID = "other" }, - func(value *store.RuntimeAllocation) { value.EnvironmentID = "other" }, + for _, mutate := range []func(*deployment.Allocation){ + func(value *deployment.Allocation) { value.TenantID = "other" }, + func(value *deployment.Allocation) { value.SessionID = "other" }, + func(value *deployment.Allocation) { value.EnvironmentID = "other" }, } { allocation := base mutate(&allocation) - resolver, err := NewResolver(resolverStore{ + resolver, err := newTestResolver(resolverStore{ session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"openai_hosted"}}`), Environment: &sessions.Environment{ID: "environment", TenantID: "tenant", SessionID: "session"}, diff --git a/services/core/internal/sandbox/providers/configuration_flow_test.go b/services/core/internal/sandbox/providers/configuration_flow_test.go index 27ad2b0ee..541b350f1 100644 --- a/services/core/internal/sandbox/providers/configuration_flow_test.go +++ b/services/core/internal/sandbox/providers/configuration_flow_test.go @@ -10,6 +10,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/api" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" @@ -93,7 +94,11 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { // registry it is built with. deployments := func() *deployment.Service { storage := deploymentpg.New(pgunit.NewPool(pool), nil) - service, err := deployment.NewService(storage, storage, registry, "") + rules, err := placement.NewRules(registry, "") + if err != nil { + t.Fatal(err) + } + service, err := deployment.NewService(storage, storage, registry, rules) if err != nil { t.Fatal(err) } @@ -147,7 +152,7 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { SessionArchive: s, Workspaces: struct{ api.EnvironmentWorkspaces }{}, }, - Sandboxes: &api.Sandboxes{Deployment: service, NodeAllocations: s, DeploymentChanges: leaseSetup{t: t, changes: changes, installation: installation}, + Sandboxes: &api.Sandboxes{Deployment: service, NodeAllocations: deploymentpg.New(pgunit.NewPool(pool), nil), DeploymentChanges: leaseSetup{t: t, changes: changes, installation: installation}, DeploymentReset: leaseSetup{t: t, changes: changes, installation: installation}, ConfigurationDiscovery: struct{ api.ConfigurationDiscovery }{}}, }) if err != nil { diff --git a/services/core/internal/store/admin_session_archive_race_test.go b/services/core/internal/store/admin_session_archive_race_test.go index 147bad1fe..8a19c4d98 100644 --- a/services/core/internal/store/admin_session_archive_race_test.go +++ b/services/core/internal/store/admin_session_archive_race_test.go @@ -31,7 +31,7 @@ func TestManagedSessionArchiveReleasesPendingNodePlacement(t *testing.T) { t.Fatal(err) } nodeID := uuid.NewString() - if _, err := nodes.Enroll(t.Context(), token, deployment.Enrollment{DeploymentGeneration: 1, SpecificationDigest: SandboxDeploymentTestSpec("docker").Digest("docker"), NodeID: nodeID, Name: "Archive fixture", Provider: "docker", Credential: strings.Repeat("x", 64), BackendFingerprint: strings.Repeat("b", 64), CoreURL: s.publicURL}); err != nil { + if _, err := nodes.Enroll(t.Context(), token, deployment.Enrollment{DeploymentGeneration: 1, SpecificationDigest: SandboxDeploymentTestSpec("docker").Digest("docker"), NodeID: nodeID, Name: "Archive fixture", Provider: "docker", Credential: strings.Repeat("x", 64), BackendFingerprint: strings.Repeat("b", 64), CoreURL: s.placement.PublicURL()}); err != nil { t.Fatal(err) } onlineManagerNode(t, s, nodeID) diff --git a/services/core/internal/store/admin_session_archive_test.go b/services/core/internal/store/admin_session_archive_test.go index 35c6d4a1e..ba8d933d1 100644 --- a/services/core/internal/store/admin_session_archive_test.go +++ b/services/core/internal/store/admin_session_archive_test.go @@ -23,7 +23,7 @@ func managedArchiveFixture(t *testing.T) (*Store, *Store, string) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) installation := uuid.NewString() changes := deploymentExecution(t, w) @@ -52,9 +52,9 @@ func managedArchiveSession(t *testing.T, s *Store, input sessions.CreateSession) return tenant, session } -func archiveAllocation(t *testing.T, w *Store, tenant string, session sessions.Session, installation string) RuntimeAllocation { +func archiveAllocation(t *testing.T, w *Store, tenant string, session sessions.Session, installation string) deployment.Allocation { t.Helper() - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, installation, runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, installation, runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } @@ -113,7 +113,7 @@ func TestManagedSessionArchiveUnallocatedAndGuards(t *testing.T) { if status, err := s.GetManagedSessionArchive(ctx, tenant, session.ID); err != nil || status != result { t.Fatal("status differs from committed archive", status, err) } - if _, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, installation, runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, installation, runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, deployment.ErrInvalidInput) { t.Fatal("archived Environment allocated after archive", err) } if _, err := s.ReserveEnvironmentInput(t.Context(), tenant, session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}); !errors.Is(err, sessions.ErrEnvironmentUnavailable) { @@ -157,24 +157,24 @@ func TestManagedSessionArchiveRetainsHistoryAndSettledResources(t *testing.T) { if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("archive retained runtime authority", err) } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), owner); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("archive discarded unknown Create ownership", err) } - replay, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, installation, runtimedevice.HashCredential(uuid.NewString())) + replay, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, installation, runtimedevice.HashCredential(uuid.NewString())) if err != nil || !replay.Replayed || replay.ID != owner.ID || replay.DeviceID != owner.DeviceID || replay.State != "cleanup_pending" { t.Fatal("late provisioning retry replaced archived allocation", replay, err) } - if _, err := w.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } current, err := s.GetSession(t.Context(), tenant, session.ID) if err != nil || current.EnvironmentFailure != nil || current.LastTurn == nil || current.LastTurn.Status != sessions.TurnCompleted || current.Environment.Status != "expired" { t.Fatal("cleanup rewrote completed outcome", current, err) } - if _, err := w.SettleRuntimeCreation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).SettleCreation(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } if result, err := s.GetManagedSessionArchive(t.Context(), tenant, session.ID); err != nil || result.State != "released" { 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 4a3f9d05e..146de71bd 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 @@ -36,6 +36,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { t.Fatal(err) } s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s.SetPlacement(fixtureRules(t, db)) installation := uuid.NewString() provider := &lifecycleProvider{resources: map[string]sandbox.Info{}} deployments := fixtureDeployment(t, db) @@ -114,7 +115,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { if w.Code != http.StatusOK || json.Unmarshal(w.Body.Bytes(), &archived) != nil || archived.State != "cleanup_pending" || archived.SessionID != active.ID { t.Fatalf("archive did not use Worker's leased Store: %d %s", w.Code, w.Body) } - allocation, err := s.GetRuntimeAllocation(t.Context(), project.TenantID, active.Environment.ID) + allocation, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: project.TenantID, EnvironmentID: active.Environment.ID}) if err != nil || allocation.ID != owner.ID || allocation.State != "cleanup_pending" { t.Fatal("archive did not retain cleanup ownership", allocation, 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 ab5d19976..67eb2f2da 100644 --- a/services/core/internal/store/agent_execution_defaults_http_test.go +++ b/services/core/internal/store/agent_execution_defaults_http_test.go @@ -20,11 +20,12 @@ func TestAgentExecutionDefaultsPublicSnapshotAndPrecedence(t *testing.T) { _, pool := store.NewTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{19}, 32)) st, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + st.SetPlacement(fixtureRules(t, db)) tenant, token := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "defaults-test", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}}) deployment := &v1.ModelProviderInput{Protocol: "responses", BaseURL: "https://deployment.example/v1", APIKey: "deployment-canary"} defaultsCalls := 0 - handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, st, db), withHarnesses([]string{"codex", "claude_sdk", "mcode"}), modelProviderDefaults(func(context.Context, string) (*modelconfiguration.Snapshot, error) { + handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, db), withHarnesses([]string{"codex", "claude_sdk", "mcode"}), modelProviderDefaults(func(context.Context, string) (*modelconfiguration.Snapshot, error) { defaultsCalls++ copy := *deployment return &modelconfiguration.Snapshot{Model: "fixture", Provider: ©, Revision: uuid.New()}, nil diff --git a/services/core/internal/store/archive_cancellation_test.go b/services/core/internal/store/archive_cancellation_test.go index 62280521a..6a34f1f5f 100644 --- a/services/core/internal/store/archive_cancellation_test.go +++ b/services/core/internal/store/archive_cancellation_test.go @@ -11,6 +11,7 @@ import ( "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/modelconfigurationpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" @@ -64,11 +65,11 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { t.Fatal(err) } secret := uuid.NewString() - owner, err := writer.ReserveRuntimeAllocation(t.Context(), project.TenantID, session.Environment.ID, installation, runtimedevice.HashCredential(secret)) + owner, err := leased.Deployment.ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: project.TenantID, EnvironmentID: session.Environment.ID}, installation, runtimedevice.HashCredential(secret)) if err != nil { t.Fatal(err) } - for _, step := range []func(context.Context, store.RuntimeAllocation) (store.RuntimeAllocation, error){writer.ObserveRuntimeRunning, writer.SettleRuntimeCreation} { + for _, step := range []func(context.Context, deployment.Allocation) (deployment.Allocation, error){leased.Deployment.ObserveRunning, leased.Deployment.SettleCreation} { owner, err = step(t.Context(), owner) if err != nil { t.Fatal(err) @@ -177,7 +178,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { if _, err := writer.ArchiveManagedSession(repeatAudit, h.tenant, session.ID, 1); err != nil { t.Fatal(err) } - if _, err := writer.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err := leased.Deployment.RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } @@ -315,7 +316,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { } // The original cleanup owner survives every delivery outcome; only // provider receipts can release its resources. - allocation, err := s.GetRuntimeAllocation(t.Context(), h.tenant, session.Environment.ID) + allocation, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: h.tenant, EnvironmentID: session.Environment.ID}) if err != nil || allocation.State != "cleanup_pending" { t.Fatal(allocation, err) } diff --git a/services/core/internal/store/archived_cancellation_migration_test.go b/services/core/internal/store/archived_cancellation_migration_test.go index d829806de..bd966700b 100644 --- a/services/core/internal/store/archived_cancellation_migration_test.go +++ b/services/core/internal/store/archived_cancellation_migration_test.go @@ -38,10 +38,10 @@ func TestArchivedCancellationMigrationDoesNotAdoptOldRevocations(t *testing.T) { if _, err := provider.Up(t.Context()); err != nil { t.Fatal(err) } - if _, err := w.SettleRuntimeCreation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).SettleCreation(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } if _, err := provider.DownTo(t.Context(), 79); err != nil { diff --git a/services/core/internal/store/credential_matrix_http_test.go b/services/core/internal/store/credential_matrix_http_test.go index 786eea348..b2621c35d 100644 --- a/services/core/internal/store/credential_matrix_http_test.go +++ b/services/core/internal/store/credential_matrix_http_test.go @@ -30,15 +30,15 @@ import ( // own /api/v1 machine connection routes. func TestCredentialNamespaceMatrix(t *testing.T) { s, db := newManagedTestStoreDB(t) - s.SetPublicURL("https://core.example") db.publicURL = "https://core.example" + s.SetPlacement(fixtureRules(t, db)) ctx := t.Context() coreKey := uuid.NewString() admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential(coreKey)}) if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, s, db), withCoreKeys(admin)) + handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, db), withCoreKeys(admin)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/deployment_fixture_test.go b/services/core/internal/store/deployment_fixture_test.go index 41b0d7139..af32eee77 100644 --- a/services/core/internal/store/deployment_fixture_test.go +++ b/services/core/internal/store/deployment_fixture_test.go @@ -4,6 +4,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" @@ -15,12 +16,31 @@ func deploymentStore(s *Store) *deploymentpg.Store { return deploymentpg.New(pgunit.NewPool(s.pool), s.credentialCipher) } +// placementRules builds the placement rules as cmd/server does, on the +// built-in providers and publicURL. +func placementRules(t testing.TB, publicURL string) *placement.Rules { + t.Helper() + rules, err := placement.NewRules(providers.Builtin(), publicURL) + if err != nil { + t.Fatal(err) + } + return rules +} + +// withPlacement gives s the placement rules cmd/server builds on the built-in +// providers and an unset public URL, and returns s. +func withPlacement(t testing.TB, s *Store) *Store { + t.Helper() + s.SetPlacement(placementRules(t, "")) + return s +} + // deploymentService builds the deployment service as cmd/server does, on s's -// database and credential key, with s's current public URL. +// database, credential key and placement rules. func deploymentService(t testing.TB, s *Store) *deployment.Service { t.Helper() adapter := deploymentStore(s) - service, err := deployment.NewService(adapter, adapter, providers.Builtin(), s.publicURL) + service, err := deployment.NewService(adapter, adapter, providers.Builtin(), s.placement) if err != nil { t.Fatal(err) } 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 66592c120..e3a09f6c6 100644 --- a/services/core/internal/store/deployment_model_providers_http_test.go +++ b/services/core/internal/store/deployment_model_providers_http_test.go @@ -36,13 +36,14 @@ func TestDeploymentModelProvidersHTTP(t *testing.T) { } cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{53}, 32)) st, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + st.SetPlacement(fixtureRules(t, db)) tenant, projectKey, coreKey := uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "defaults-http", TokenSHA256: runtimedevice.HashCredential(projectKey), TenantID: tenant}}) admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential(coreKey)}) if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, st, db), withCoreKeys(admin), withHarnesses([]string{"codex", "mcode"})) + handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, db), withCoreKeys(admin), withHarnesses([]string{"codex", "mcode"})) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/deployment_public_fixture_test.go b/services/core/internal/store/deployment_public_fixture_test.go index 7a041cc26..f7377971d 100644 --- a/services/core/internal/store/deployment_public_fixture_test.go +++ b/services/core/internal/store/deployment_public_fixture_test.go @@ -4,6 +4,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" @@ -20,8 +21,28 @@ func fixtureDeployment(t testing.TB, db fixtureDB) *deployment.Service { } func fixtureDeploymentService(db fixtureDB) (*deployment.Service, error) { + rules, err := placement.NewRules(providers.Builtin(), db.publicURL) + if err != nil { + return nil, err + } adapter := deploymentpg.New(pgunit.NewPool(db.pool), db.cipher) - return deployment.NewService(adapter, adapter, providers.Builtin(), db.publicURL) + return deployment.NewService(adapter, adapter, providers.Builtin(), rules) +} + +// fixtureRules builds the placement rules on db's public URL, as cmd/server +// does for the Store's Session creation. +func fixtureRules(t testing.TB, db fixtureDB) *placement.Rules { + t.Helper() + rules, err := placement.NewRules(providers.Builtin(), db.publicURL) + if err != nil { + t.Fatal(err) + } + return rules +} + +// fixtureReader is the deployment reader on db, as cmd/server builds it. +func fixtureReader(db fixtureDB) *deploymentpg.Store { + return deploymentpg.New(pgunit.NewPool(db.pool), db.cipher) } // fixtureOwnerEpoch reads the execution owner epoch from the deployment store, diff --git a/services/core/internal/store/device_bootstrap_binding_test.go b/services/core/internal/store/device_bootstrap_binding_test.go index d0c94f706..d01d32146 100644 --- a/services/core/internal/store/device_bootstrap_binding_test.go +++ b/services/core/internal/store/device_bootstrap_binding_test.go @@ -20,7 +20,7 @@ func TestDeviceCredentialCarriesPersistedAllocationNode(t *testing.T) { } remote := uuid.NewString() _, err = nodes.Enroll(t.Context(), token, deployment.Enrollment{DeploymentGeneration: 1, SpecificationDigest: SandboxDeploymentTestSpec("docker").Digest("docker"), NodeID: remote, Credential: strings.Repeat("x", 64), - Name: "remote", Provider: "docker", BackendFingerprint: strings.Repeat("b", 64), CoreURL: s.publicURL}) + Name: "remote", Provider: "docker", BackendFingerprint: strings.Repeat("b", 64), CoreURL: s.placement.PublicURL()}) if err != nil { t.Fatal(err) } @@ -36,7 +36,7 @@ func TestDeviceCredentialCarriesPersistedAllocationNode(t *testing.T) { if err != nil { t.Fatal(err) } - allocation, err := writer.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, d.InstallationID, runtimedevice.HashCredential(bearer)) + allocation, err := deploymentExecution(t, writer).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, d.InstallationID, runtimedevice.HashCredential(bearer)) if err != nil { t.Fatal(err) } @@ -66,7 +66,7 @@ func TestDeviceCredentialWithoutManagedNodeRetainsPublicRouteIdentity(t *testing t.Fatal(err) } _, environment := localEnvironment(t, s, tenant) - allocation, err := executionWriter(t, s).ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, uuid.NewString(), runtimedevice.HashCredential("allocation-token")) + allocation, err := deploymentExecution(t, executionWriter(t, s)).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, uuid.NewString(), runtimedevice.HashCredential("allocation-token")) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_connections.go b/services/core/internal/store/environment_connections.go deleted file mode 100644 index 6b95a4e72..000000000 --- a/services/core/internal/store/environment_connections.go +++ /dev/null @@ -1,14 +0,0 @@ -package store - -import ( - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5/pgtype" -) - -func parseConnectionGeneration(value string) (pgtype.UUID, error) { - id, err := parseID(value) - if err != nil || id.Bytes == [16]byte{} { - return pgtype.UUID{}, sessions.ErrInvalidInput - } - return id, nil -} diff --git a/services/core/internal/store/environment_initialization_test.go b/services/core/internal/store/environment_initialization_test.go index 5504ef831..05b4be525 100644 --- a/services/core/internal/store/environment_initialization_test.go +++ b/services/core/internal/store/environment_initialization_test.go @@ -14,6 +14,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "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" @@ -134,7 +135,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test want = "failed" } awaitInitialization(t, db.pool, principal.TenantID, environment.ID, want) - if _, err := s.GetRuntimeAllocation(t.Context(), principal.TenantID, environment.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: principal.TenantID, EnvironmentID: environment.ID}); !errors.Is(err, deployment.ErrNotFound) { t.Fatal("self-hosted preparation fabricated allocation", err) } if outcome == "completed" { diff --git a/services/core/internal/store/environment_plugins_test.go b/services/core/internal/store/environment_plugins_test.go index 241af01d6..1a0ffa94f 100644 --- a/services/core/internal/store/environment_plugins_test.go +++ b/services/core/internal/store/environment_plugins_test.go @@ -21,7 +21,7 @@ func TestPluginsFrozenInSession(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) var archive bytes.Buffer writer := zip.NewWriter(&archive) for path, body := range map[string]string{ diff --git a/services/core/internal/store/environment_setup_test.go b/services/core/internal/store/environment_setup_test.go index a7e9b4134..04c5e9456 100644 --- a/services/core/internal/store/environment_setup_test.go +++ b/services/core/internal/store/environment_setup_test.go @@ -20,7 +20,7 @@ func TestEnvironmentSetupEncryptedSnapshotAndIsolation(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) tenant, foreign := uuid.NewString(), uuid.NewString() setup := environmentconfig.Setup{Env: map[string]string{"SECRET": "session-env-canary"}, Commands: []environmentconfig.SetupCommand{{Command: "printf session-command-canary > result"}}, Packages: v1.EnvironmentPackages{NPM: []string{"is-number@7.0.0"}}} input := sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"environment":{"type":"openai_hosted"}}`), Initialization: setup} diff --git a/services/core/internal/store/environment_skill_references_test.go b/services/core/internal/store/environment_skill_references_test.go index d57c021fd..a0820cf48 100644 --- a/services/core/internal/store/environment_skill_references_test.go +++ b/services/core/internal/store/environment_skill_references_test.go @@ -20,7 +20,7 @@ func TestSkillReferencesFreezeWithinSessionCreation(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) skillService := SkillService(t, pool, cipher) tenant := uuid.NewString() first, second := skillArchive(t, "frozen-first"), skillArchive(t, "frozen-second") @@ -130,7 +130,7 @@ func TestSkillReferenceAuthorizationRollsBackSession(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) tenant, foreign := uuid.NewString(), uuid.NewString() skill, err := SkillService(t, pool, cipher).CreateSkill(t.Context(), skills.CreateSkill{TenantID: tenant, Archive: skillArchive(t, "private-owner")}) if err != nil { diff --git a/services/core/internal/store/environment_skills_test.go b/services/core/internal/store/environment_skills_test.go index 52a689a0c..8e25d5743 100644 --- a/services/core/internal/store/environment_skills_test.go +++ b/services/core/internal/store/environment_skills_test.go @@ -20,7 +20,7 @@ func TestSkillsFrozenInSession(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) var archive bytes.Buffer writer := zip.NewWriter(&archive) header := &zip.FileHeader{Name: "proof/SKILL.md", Method: zip.Store} diff --git a/services/core/internal/store/execution.go b/services/core/internal/store/execution.go index 05238c038..0978511c5 100644 --- a/services/core/internal/store/execution.go +++ b/services/core/internal/store/execution.go @@ -1,7 +1,6 @@ package store import ( - "context" "errors" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" @@ -28,12 +27,3 @@ func (s *Store) checkExecutionAuthority() error { } return nil } - -// 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) -} diff --git a/services/core/internal/store/execution_test.go b/services/core/internal/store/execution_test.go index 1ccecaa1e..4b0c2f5ba 100644 --- a/services/core/internal/store/execution_test.go +++ b/services/core/internal/store/execution_test.go @@ -262,9 +262,8 @@ func TestPooledStoreHasNoExecutionAuthority(t *testing.T) { _, archiveErr := s.ArchiveManagedSession(t.Context(), tenant, session.ID, 0) _, expiryErr := s.ExpireEnvironmentInputs(t.Context()) for name, err := range map[string]error{ - "ownership check": s.checkExecutionOwnership(t.Context()), - "archive": archiveErr, - "input expiry": expiryErr, + "archive": archiveErr, + "input expiry": expiryErr, } { if !errors.Is(err, ErrExecutionAuthority) { t.Fatalf("pooled Store ran %s: %v", name, err) diff --git a/services/core/internal/store/export_test.go b/services/core/internal/store/export_test.go index fc18300a2..89ad806ba 100644 --- a/services/core/internal/store/export_test.go +++ b/services/core/internal/store/export_test.go @@ -31,7 +31,7 @@ func FixtureCipher() *credentialcrypto.Cipher { return fixtureCipher } // can freeze a model provider. func NewModelTestStore(t *testing.T) (*Store, *pgxpool.Pool) { _, pool := testStore(t) - return NewWithCredentialCipher(pool, fixtureCipher), pool + return withPlacement(t, NewWithCredentialCipher(pool, fixtureCipher)), pool } // FixtureModelProvider is a valid bundle for the harness. Hosted and self-hosted diff --git a/services/core/internal/store/fixture_db_test.go b/services/core/internal/store/fixture_db_test.go index 88b605717..3372f2b31 100644 --- a/services/core/internal/store/fixture_db_test.go +++ b/services/core/internal/store/fixture_db_test.go @@ -22,7 +22,7 @@ import ( type fixtureDB struct { pool *pgxpool.Pool cipher *credentialcrypto.Cipher // nil for a keyless Store - publicURL string // the value given to the Store's SetPublicURL, if any + publicURL string // the public URL of the Store's placement rules, if any } // newTestStoreDB is store.NewTestStore with the fixtureDB that built it. 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 6ea80d866..ebb1dd329 100644 --- a/services/core/internal/store/hosted_initialization_failure_public_test.go +++ b/services/core/internal/store/hosted_initialization_failure_public_test.go @@ -106,7 +106,9 @@ func hostedFailureStore(t *testing.T) (*store.Store, fixtureDB) { if err != nil { t.Fatal(err) } - return store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s.SetPlacement(fixtureRules(t, db)) + return s, db } func hostedFailureSession(t *testing.T, s *store.Store, db fixtureDB, tenant string, input sessions.CreateSession) (sessions.Session, sessions.Environment) { diff --git a/services/core/internal/store/initial_files_http_test.go b/services/core/internal/store/initial_files_http_test.go index beb136f88..8dda7ceef 100644 --- a/services/core/internal/store/initial_files_http_test.go +++ b/services/core/internal/store/initial_files_http_test.go @@ -21,10 +21,11 @@ func TestInitialFilesHTTPInlineLimitsAndRetry(t *testing.T) { t.Fatal(err) } s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s.SetPlacement(fixtureRules(t, db)) tenant, token := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}}) // Exercise HTTP parsing and durable storage without starting a Runtime. - handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), managedSandboxes(t, s, db), fixtureDeploymentProvider()) + handler, err := publicHandler(t, s, db, auth, "codex", storeExecution(t, s), managedSandboxes(t, db), fixtureDeploymentProvider()) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/initial_files_test.go b/services/core/internal/store/initial_files_test.go index aaf637331..0bc20c97d 100644 --- a/services/core/internal/store/initial_files_test.go +++ b/services/core/internal/store/initial_files_test.go @@ -42,7 +42,7 @@ func TestInitialFilesFrozenEncryptedIsolatedAndRetryable(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) tenant, foreign := uuid.NewString(), uuid.NewString() canary := []byte("private-initial-file-canary\x00\xff") sourceFiles, _ := testFiles(t, pool) diff --git a/services/core/internal/store/list_cursor_public_test.go b/services/core/internal/store/list_cursor_public_test.go index 8ad3e30bd..009578904 100644 --- a/services/core/internal/store/list_cursor_public_test.go +++ b/services/core/internal/store/list_cursor_public_test.go @@ -227,6 +227,7 @@ func TestListCursorErrorsPostgres(t *testing.T) { t.Fatal(err) } s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s.SetPlacement(fixtureRules(t, db)) owner, foreign := uuid.NewString(), uuid.NewString() ownerTenant, foreignTenant := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ diff --git a/services/core/internal/store/managed_test_store_test.go b/services/core/internal/store/managed_test_store_test.go index b6c2db27f..172392a82 100644 --- a/services/core/internal/store/managed_test_store_test.go +++ b/services/core/internal/store/managed_test_store_test.go @@ -14,7 +14,7 @@ func newManagedTestStore(t *testing.T) (*Store, *pgxpool.Pool) { t.Helper() pool := pgtest.OpenIsolated(t, nil) // Hosted Sessions freeze a model provider, which needs a credential key. - return NewWithCredentialCipher(pool, fixtureCipher), pool + return withPlacement(t, NewWithCredentialCipher(pool, fixtureCipher)), pool } func NewManagedTestStore(t *testing.T) (*Store, *pgxpool.Pool) { diff --git a/services/core/internal/store/public_handler_fixture_test.go b/services/core/internal/store/public_handler_fixture_test.go index 923187664..82da7f7e6 100644 --- a/services/core/internal/store/public_handler_fixture_test.go +++ b/services/core/internal/store/public_handler_fixture_test.go @@ -181,11 +181,11 @@ func executorURL(url string) func(*api.Dependencies) { // administration and node routes and openai_hosted Environments. Deployment // changes, reset and discovery need the Worker and are strict stand-ins. It // follows the option that enables Execution. -func managedSandboxes(t testing.TB, s *store.Store, db fixtureDB) func(*api.Dependencies) { +func managedSandboxes(t testing.TB, db fixtureDB) func(*api.Dependencies) { return func(d *api.Dependencies) { d.Sandboxes = &api.Sandboxes{ Deployment: fixtureDeployment(t, db), - NodeAllocations: s, + NodeAllocations: fixtureReader(db), DeploymentChanges: strictStandIn{t}, DeploymentReset: strictStandIn{t}, ConfigurationDiscovery: strictStandIn{t}, diff --git a/services/core/internal/store/public_url.go b/services/core/internal/store/public_url.go deleted file mode 100644 index d3a6ee70c..000000000 --- a/services/core/internal/store/public_url.go +++ /dev/null @@ -1,5 +0,0 @@ -package store - -// SetPublicURL records OAC_PUBLIC_URL, validated by the caller. Placement -// admits only nodes enrolled with it. Call it once, before serving requests. -func (s *Store) SetPublicURL(value string) { s.publicURL = value } diff --git a/services/core/internal/store/runtime_adoption_test.go b/services/core/internal/store/runtime_adoption_test.go index ce25b5caf..dea4dcb4c 100644 --- a/services/core/internal/store/runtime_adoption_test.go +++ b/services/core/internal/store/runtime_adoption_test.go @@ -10,7 +10,7 @@ import ( "github.com/google/uuid" ) -func legacyAdoptionFixture(t *testing.T) (*Store, *Store, deployment.ProcessDeployment, RuntimeAllocation) { +func nodelessAllocationFixture(t *testing.T) (*Store, *Store, deployment.ProcessDeployment, deployment.Allocation) { t.Helper() s, _ := newManagedTestStore(t) w := executionWriter(t, s) @@ -18,7 +18,7 @@ func legacyAdoptionFixture(t *testing.T) (*Store, *Store, deployment.ProcessDepl deploymentConfigure(t, w, &d) tenant := uuid.NewString() _, e := localEnvironment(t, s, tenant) - a, err := w.ReserveRuntimeAllocation(t.Context(), tenant, e.ID, d.InstallationID, runtimedevice.HashCredential("runtime")) + a, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: e.ID}, d.InstallationID, runtimedevice.HashCredential("runtime")) if err != nil { t.Fatal(err) } @@ -28,7 +28,7 @@ func legacyAdoptionFixture(t *testing.T) (*Store, *Store, deployment.ProcessDepl d.LocalMaxActive, d.LocalMaxRetained = 4, 16 return s, w, d, a } -func requireNoLegacyBinding(t *testing.T, s *Store) { +func requireNoNodeBinding(t *testing.T, s *Store) { t.Helper() var nodes, placements, bound int var kind string @@ -43,12 +43,12 @@ func TestHistoricalRuntimeResourcesCannotBeAdopted(t *testing.T) { for _, state := range []string{"creating", "running", "cleanup_pending", "released"} { for _, pending := range []bool{false, true} { t.Run(fmt.Sprintf("%s/pending=%t", state, pending), func(t *testing.T) { - s, w, d, a := legacyAdoptionFixture(t) + s, w, d, a := nodelessAllocationFixture(t) runtimeSuspensionSQL(t, s.pool, `UPDATE runtime_allocations SET state=$2,create_settled=($2='released'),released_at=CASE WHEN $2='released' THEN clock_timestamp() ELSE NULL END WHERE id=$1`, a.ID, state) if pending { _, _ = localEnvironment(t, s, a.TenantID) } - before, err := s.GetRuntimeAllocation(t.Context(), a.TenantID, a.EnvironmentID) + before, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: a.TenantID, EnvironmentID: a.EnvironmentID}) if err != nil { t.Fatal(err) } @@ -61,9 +61,9 @@ func TestHistoricalRuntimeResourcesCannotBeAdopted(t *testing.T) { if err == nil { t.Fatal("historical resources adopted") } - requireNoLegacyBinding(t, s) + requireNoNodeBinding(t, s) } - after, err := s.GetRuntimeAllocation(t.Context(), a.TenantID, a.EnvironmentID) + after, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: a.TenantID, EnvironmentID: a.EnvironmentID}) if err != nil || !reflect.DeepEqual(before, after) { t.Fatal("retained receipt changed", err) } diff --git a/services/core/internal/store/runtime_allocation_state.go b/services/core/internal/store/runtime_allocation_state.go deleted file mode 100644 index 669e9446e..000000000 --- a/services/core/internal/store/runtime_allocation_state.go +++ /dev/null @@ -1,157 +0,0 @@ -package store - -import ( - "context" - "errors" - - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -// ObserveRuntimeRunning requires verified original compute identity. It does not -// mark the Environment connected or qualify native preparation. -func (s *Store) ObserveRuntimeRunning(ctx context.Context, owner RuntimeAllocation) (RuntimeAllocation, error) { - return s.mutateRuntimeAllocation(ctx, owner, true, func(ctx context.Context, q *sqlc.Queries, row sqlc.RuntimeAllocation) (sqlc.RuntimeAllocation, error) { - return q.ObserveRuntimeRunning(ctx, row.ID) - }) -} - -// KeepRuntimeAllocation follows an authenticated connection and successful -// provider observation. A keepalive cannot revive cleanup or an expired lease. -func (s *Store) KeepRuntimeAllocation(ctx context.Context, owner RuntimeAllocation) (RuntimeAllocation, error) { - return s.mutateRuntimeAllocation(ctx, owner, true, func(ctx context.Context, q *sqlc.Queries, row sqlc.RuntimeAllocation) (sqlc.RuntimeAllocation, error) { - return q.KeepRuntimeAllocation(ctx, row.ID) - }) -} - -// SettleRuntimeCreation records evidence that the original Create can no longer -// mutate resources. A timeout, missing container or lost lease is not evidence. -func (s *Store) SettleRuntimeCreation(ctx context.Context, owner RuntimeAllocation) (RuntimeAllocation, error) { - return s.mutateRuntimeAllocation(ctx, owner, false, func(ctx context.Context, q *sqlc.Queries, row sqlc.RuntimeAllocation) (sqlc.RuntimeAllocation, error) { - if row.State == "released" { - return row, nil - } - return q.SettleRuntimeCreation(ctx, row.ID) - }) -} - -// RequestRuntimeCleanup revokes future authority before external reclamation. -// Cancellation requests do not prove existing native work has stopped. A live, -// unexpired Environment fails with the generic provisioning reason. -func (s *Store) RequestRuntimeCleanup(ctx context.Context, owner RuntimeAllocation) (RuntimeAllocation, error) { - return s.requestRuntimeCleanup(ctx, owner, sessions.ProvisioningFailureReason, nil, false) -} - -// ReleaseAbsentRuntimeCreation consumes provider proof that the original attempt -// is settled and owns no resources. Authority revocation and release commit together. -func (s *Store) ReleaseAbsentRuntimeCreation(ctx context.Context, owner RuntimeAllocation) (RuntimeAllocation, error) { - return s.requestRuntimeCleanup(ctx, owner, sessions.ProvisioningFailureReason, nil, true) -} - -func (s *Store) requestRuntimeCleanup(ctx context.Context, owner RuntimeAllocation, reason string, detail *sessions.ProvisioningFailureDetail, absent bool) (RuntimeAllocation, error) { - return s.mutateRuntimeAllocation(ctx, owner, false, func(ctx context.Context, q *sqlc.Queries, row sqlc.RuntimeAllocation) (sqlc.RuntimeAllocation, error) { - if row.State == "released" { - return row, nil - } - tenant, _ := parseID(owner.TenantID) - if _, err := q.RevokeRuntimeCleanupDevice(ctx, sqlc.RevokeRuntimeCleanupDeviceParams{TenantID: tenant, DeviceID: row.DeviceID}); err != nil { - return sqlc.RuntimeAllocation{}, err - } - current, err := q.GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: tenant, EnvironmentID: row.EnvironmentID}) - if err != nil { - return sqlc.RuntimeAllocation{}, err - } - session := sessionpg.BindSession(q, current.TenantID, current.SessionID) - if current.DeletedAt.Valid { - err = sessions.CancelWork(ctx, session) - } else { - err = sessions.TerminateEnvironment(ctx, session, current.Expired, reason, detail) - } - if err != nil { - return sqlc.RuntimeAllocation{}, err - } - - pending, err := q.RequestRuntimeCleanup(ctx, row.ID) - if err != nil || !absent { - return pending, err - } - if _, err := q.SettleRuntimeCreation(ctx, row.ID); err != nil { - return sqlc.RuntimeAllocation{}, err - } - released, err := q.ReleaseRuntimeAllocation(ctx, row.ID) - if err != nil { - return sqlc.RuntimeAllocation{}, err - } - if err := q.ReleaseRuntimePlacement(ctx, row.EnvironmentID); err != nil { - return sqlc.RuntimeAllocation{}, err - } - return released, nil - }) -} - -// ReleaseRuntimeAllocation follows successful owned container/volume cleanup. -// Unknown Create outcomes keep their cleanup record even when compute is absent. -func (s *Store) ReleaseRuntimeAllocation(ctx context.Context, owner RuntimeAllocation) (RuntimeAllocation, error) { - return s.mutateRuntimeAllocation(ctx, owner, false, func(ctx context.Context, q *sqlc.Queries, row sqlc.RuntimeAllocation) (sqlc.RuntimeAllocation, error) { - if row.State == "released" { - return row, nil - } - released, err := q.ReleaseRuntimeAllocation(ctx, row.ID) - if err != nil { - return released, err - } - if err := q.ReleaseRuntimePlacement(ctx, row.EnvironmentID); err != nil { - return sqlc.RuntimeAllocation{}, err - } - return released, nil - }) -} - -func (s *Store) mutateRuntimeAllocation(ctx context.Context, owner RuntimeAllocation, live bool, apply func(context.Context, *sqlc.Queries, sqlc.RuntimeAllocation) (sqlc.RuntimeAllocation, error)) (RuntimeAllocation, error) { - if err := s.checkExecutionAuthority(); err != nil { - return RuntimeAllocation{}, err - } - previous, err := s.GetRuntimeAllocation(ctx, owner.TenantID, owner.EnvironmentID) - if err != nil { - return RuntimeAllocation{}, err - } - if previous.ID != owner.ID || previous.DeviceID != owner.DeviceID || previous.ProviderKey != owner.ProviderKey || previous.NodeID != owner.NodeID { - return RuntimeAllocation{}, sessions.ErrIdempotencyConflict - } - lookup, _ := sessionpg.DeviceLookup(owner.TenantID, owner.EnvironmentID) - var result RuntimeAllocation - err = s.withSession(ctx, owner.TenantID, previous.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - current, err := q.GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: lookup.TenantID, EnvironmentID: lookup.ID}) - if err != nil { - return err - } - if live { - if current.DeletedAt.Valid { - return sessions.ErrNotFound - } - device, err := q.GetSessionDevice(ctx, sqlc.GetSessionDeviceParams{TenantID: lookup.TenantID, ID: session}) - if err != nil { - return err - } - if device.ID != current.RuntimeAllocation.DeviceID || device.EnvironmentID != lookup.ID { - return sessions.ErrDeviceBindingConflict - } - } - row, err := apply(ctx, q, current.RuntimeAllocation) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrTurnConflict - } - if err == nil { - result = runtimeAllocationFromRow(row, session, lookup.TenantID, current.DeletedAt, current.Expired) - } - return err - }) - if err != nil { - return RuntimeAllocation{}, err - } - return result, nil -} diff --git a/services/core/internal/store/runtime_allocations.go b/services/core/internal/store/runtime_allocations.go deleted file mode 100644 index 1979934cc..000000000 --- a/services/core/internal/store/runtime_allocations.go +++ /dev/null @@ -1,256 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "time" - - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -// RuntimeAllocation retains compute ownership, not public readiness. It survives -// Session deletion until cleanup is confirmed. No bootstrap secret is retained. -type RuntimeAllocation struct { - DeploymentGeneration uint64 - NodeID string - ObservationError string - ComputePhase string - ComputeRevision int64 - ComputeState json.RawMessage - ComputeActivityAt time.Time - ComputeWakeRequested bool - ComputeRetainedUntil *time.Time - ID, EnvironmentID, SessionID, TenantID, DeviceID, ProviderKey string - State string - CreateSettled, SessionDeleted, Replayed, Expired bool - CreatedAt, KeptAt time.Time -} - -// RuntimeObservationSession is the minimum durable Core identity needed by the -// deployment-wide read-only sampler. Provider identity is resolved again by the -// observation service before any external read. -type RuntimeObservationSession struct { - TenantID, SessionID string -} - -type RuntimeObservationSessionPage struct { - Sessions []RuntimeObservationSession - NextCursor string -} - -// ReserveRuntimeAllocation commits the allocation and dedicated device together -// before external Create. Only a fresh receipt authorizes that one Create call. -func (s *Store) ReserveRuntimeAllocation(ctx context.Context, tenant, environment, providerKey, credentialHash string) (RuntimeAllocation, error) { - if err := s.checkExecutionAuthority(); err != nil { - return RuntimeAllocation{}, err - } - provider, err := parseConnectionGeneration(providerKey) - if err != nil { - return RuntimeAllocation{}, err - } - owned, err := sessionpg.LoadEnvironment(ctx, s.queries, tenant, environment) - if err != nil { - return RuntimeAllocation{}, err - } - var config struct { - Type string `json:"type"` - } - if json.Unmarshal(owned.Configuration, &config) != nil || config.Type != "openai_hosted" { - return RuntimeAllocation{}, sessions.ErrInvalidInput - } - lookup, err := sessionpg.DeviceLookup(tenant, environment) - if err != nil { - return RuntimeAllocation{}, err - } - registration, err := sessions.NewDeviceRegistration("managed-runtime", credentialHash) - if err != nil { - return RuntimeAllocation{}, err - } - deviceID := pgtype.UUID{Bytes: uuid.New(), Valid: true} - var result RuntimeAllocation - err = s.withPublicSession(ctx, tenant, owned.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - previous, err := q.GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: lookup.TenantID, EnvironmentID: lookup.ID}) - if err == nil { - if previous.RuntimeAllocation.ProviderKey != provider { - return sessions.ErrIdempotencyConflict - } - result = runtimeAllocationFromRow(previous.RuntimeAllocation, session, lookup.TenantID, previous.DeletedAt, previous.Expired) - result.Replayed = true - return nil - } - if !errors.Is(err, pgx.ErrNoRows) { - return err - } - if err := checkRuntimeDeploymentAdmission(ctx, q, providerKey); err != nil { - return err - } - current, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: lookup.TenantID, ID: session}) - if err != nil { - return err - } - if current.Environment.Status == "failed" || current.Environment.Status == "expired" { - return sessions.ErrInvalidInput - } - var nodeID pgtype.UUID - active, err := q.GetRuntimeDeployment(ctx) - if err != nil { - return err - } - generation := pgtype.Int8{Int64: active.Generation, Valid: true} - if active.Mode == "nodes" { - placement, err := q.GetRuntimePlacement(ctx, lookup.ID) - if err != nil { - return err - } - if placement.ReleasedAt.Valid || !placement.Available { - return deployment.ErrNodeUnavailable - } - nodeID = placement.NodeID - generation = placement.DeploymentGeneration - } - dedicated := sessions.ExecutionDevice{ID: uuid.UUID(deviceID.Bytes).String(), Name: registration.Name, EnvironmentID: owned.ID} - if err := sessions.CreateEnvironmentDevice(ctx, sessionpg.BindSession(q, lookup.TenantID, session), dedicated, registration.CredentialHash); err != nil { - return err - } - row, err := q.CreateRuntimeAllocation(ctx, sqlc.CreateRuntimeAllocationParams{ - ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, EnvironmentID: lookup.ID, - DeviceID: deviceID, ProviderKey: provider, NodeID: nodeID, DeploymentGeneration: generation, - }) - if err == nil { - result = runtimeAllocationFromRow(row, session, lookup.TenantID, pgtype.Timestamptz{}, false) - } - return err - }) - if err != nil { - return RuntimeAllocation{}, err - } - return result, nil -} - -// GetRuntimeAllocation is an internal cleanup lookup, including deleted Sessions. -func (s *Store) GetRuntimeAllocation(ctx context.Context, tenant, environment string) (RuntimeAllocation, error) { - lookup, err := sessionpg.DeviceLookup(tenant, environment) - if err != nil { - return RuntimeAllocation{}, err - } - row, err := s.queries.GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: lookup.TenantID, EnvironmentID: lookup.ID}) - if errors.Is(err, pgx.ErrNoRows) { - return RuntimeAllocation{}, sessions.ErrNotFound - } - if err != nil { - return RuntimeAllocation{}, err - } - return runtimeAllocationFromRow(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired), nil -} - -// 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 { - return nil, err - } - id := pgtype.UUID{Valid: true} - if after != "" { - var err error - id, err = parseID(after) - if err != nil { - return nil, err - } - } - rows, err := s.queries.ListRuntimeAllocations(ctx, id) - if err != nil { - return nil, err - } - result := make([]RuntimeAllocation, 0, len(rows)) - for _, row := range rows { - result = append(result, runtimeAllocationFromRow(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired)) - } - return result, nil -} - -// ListRuntimeObservationSessions performs a deployment-wide, read-only keyset -// scan of live managed Session identities. It excludes deleted Sessions and -// released allocations; it does not acquire, renew, or mutate Runtime state. -func (s *Store) ListRuntimeObservationSessions(ctx context.Context, after string, limit int) (RuntimeObservationSessionPage, error) { - if limit < 1 || limit > 100 { - return RuntimeObservationSessionPage{}, sessions.ErrInvalidInput - } - id := pgtype.UUID{Valid: true} - if after != "" { - var err error - id, err = parseID(after) - if err != nil { - return RuntimeObservationSessionPage{}, err - } - } - rows, err := s.queries.ListRuntimeObservationSessions(ctx, sqlc.ListRuntimeObservationSessionsParams{ID: id, Limit: int32(limit + 1)}) - if err != nil { - return RuntimeObservationSessionPage{}, err - } - page := RuntimeObservationSessionPage{Sessions: make([]RuntimeObservationSession, 0, min(limit, len(rows)))} - if len(rows) > limit { - page.NextCursor = uuid.UUID(rows[limit-1].ID.Bytes).String() - rows = rows[:limit] - } - for _, row := range rows { - page.Sessions = append(page.Sessions, RuntimeObservationSession{ - TenantID: uuid.UUID(row.TenantID.Bytes).String(), - SessionID: uuid.UUID(row.ID.Bytes).String(), - }) - } - return page, nil -} - -func runtimeAllocationFromRow(row sqlc.RuntimeAllocation, session, tenant pgtype.UUID, deleted pgtype.Timestamptz, expired bool) RuntimeAllocation { - var retainedUntil *time.Time - if row.ComputeRetainedUntil.Valid { - value := row.ComputeRetainedUntil.Time - retainedUntil = &value - } - return RuntimeAllocation{ - DeploymentGeneration: uint64(row.DeploymentGeneration.Int64), NodeID: runtimeUUID(row.NodeID), ObservationError: row.ObservationError, - ComputePhase: row.ComputePhase, ComputeRevision: row.ComputeRevision, ComputeState: row.ComputeState, - ComputeActivityAt: row.ComputeActivityAt.Time, ComputeWakeRequested: row.ComputeWakeRequested, ComputeRetainedUntil: retainedUntil, - ID: uuid.UUID(row.ID.Bytes).String(), EnvironmentID: uuid.UUID(row.EnvironmentID.Bytes).String(), - SessionID: uuid.UUID(session.Bytes).String(), TenantID: uuid.UUID(tenant.Bytes).String(), - DeviceID: uuid.UUID(row.DeviceID.Bytes).String(), ProviderKey: uuid.UUID(row.ProviderKey.Bytes).String(), - State: row.State, CreateSettled: row.CreateSettled, SessionDeleted: deleted.Valid, Expired: expired, - CreatedAt: row.CreatedAt.Time, KeptAt: row.KeptAt.Time, - } -} - -// UnallocatedHostedEnvironment is a committed resource awaiting service bootstrap. -// A missing allocation is distinct from an unknown outcome of an existing Create. -type UnallocatedHostedEnvironment struct { - ID, TenantID string -} - -func (s *Store) ListUnallocatedHostedEnvironments(ctx context.Context, after string) ([]UnallocatedHostedEnvironment, error) { - if err := s.checkExecutionOwnership(ctx); err != nil { - return nil, err - } - id := pgtype.UUID{Valid: true} - if after != "" { - var err error - id, err = parseID(after) - if err != nil { - return nil, err - } - } - rows, err := s.queries.ListUnallocatedHostedEnvironments(ctx, id) - if err != nil { - return nil, err - } - result := make([]UnallocatedHostedEnvironment, 0, len(rows)) - for _, row := range rows { - result = append(result, UnallocatedHostedEnvironment{ID: uuid.UUID(row.ID.Bytes).String(), TenantID: uuid.UUID(row.TenantID.Bytes).String()}) - } - return result, nil -} diff --git a/services/core/internal/store/runtime_allocations_test.go b/services/core/internal/store/runtime_allocations_test.go index 1e2af94d5..9ee59b425 100644 --- a/services/core/internal/store/runtime_allocations_test.go +++ b/services/core/internal/store/runtime_allocations_test.go @@ -7,6 +7,7 @@ import ( "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -19,7 +20,7 @@ func TestRuntimeAllocationAtomicOwnershipAndRecovery(t *testing.T) { tenant, provider := uuid.NewString(), uuid.NewString() session, environment := localEnvironment(t, s, tenant) secret := uuid.NewString() - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, provider, runtimedevice.HashCredential(secret)) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, provider, runtimedevice.HashCredential(secret)) if err != nil || owner.Replayed || owner.State != "creating" || owner.CreateSettled { t.Fatalf("reservation: %+v %v", owner, err) } @@ -27,10 +28,10 @@ func TestRuntimeAllocationAtomicOwnershipAndRecovery(t *testing.T) { if err != nil || bound.ID != owner.DeviceID || bound.EnvironmentID != environment.ID { t.Fatalf("binding not committed with allocation: %+v %v", bound, err) } - if _, err := w.ReserveRuntimeAllocation(t.Context(), uuid.NewString(), environment.ID, provider, runtimedevice.HashCredential(secret)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: uuid.NewString(), EnvironmentID: environment.ID}, provider, runtimedevice.HashCredential(secret)); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("foreign allocation accepted: %v", err) } - if _, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, uuid.NewString(), runtimedevice.HashCredential(secret)); !errors.Is(err, sessions.ErrIdempotencyConflict) { + if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, uuid.NewString(), runtimedevice.HashCredential(secret)); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatalf("provider target changed: %v", err) } awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, w.pool) @@ -38,12 +39,12 @@ func TestRuntimeAllocationAtomicOwnershipAndRecovery(t *testing.T) { t.Fatal(err) } awaitRelease() - if _, err := w.ObserveRuntimeRunning(t.Context(), owner); err == nil { + if _, err := deploymentExecution(t, w).ObserveRunning(t.Context(), owner); err == nil { t.Fatal("lost writer changed allocation") } reopened, _ := testStore(t) next := executionWriter(t, reopened) - retry, err := next.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, provider, runtimedevice.HashCredential(uuid.NewString())) + retry, err := deploymentExecution(t, next).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, provider, runtimedevice.HashCredential(uuid.NewString())) if err != nil || !retry.Replayed || retry.ID != owner.ID || retry.DeviceID != owner.DeviceID { t.Fatalf("restart replaced unknown allocation: %+v %v", retry, err) } @@ -54,22 +55,22 @@ func TestRuntimeAllocationAtomicOwnershipAndRecovery(t *testing.T) { if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { t.Fatal(err) } - retained, err := next.GetRuntimeAllocation(t.Context(), tenant, environment.ID) + retained, err := deploymentStore(next).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || !retained.SessionDeleted || retained.ID != owner.ID { t.Fatalf("deletion discarded cleanup identity: %+v %v", retained, err) } - if _, err := next.ObserveRuntimeRunning(t.Context(), owner); !errors.Is(err, sessions.ErrNotFound) { + if _, err := deploymentExecution(t, next).ObserveRunning(t.Context(), owner); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("late creation revived deleted Session: %v", err) } - if _, err := next.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, next).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err := next.ReleaseRuntimeAllocation(t.Context(), owner); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, next).ReleaseAllocation(t.Context(), owner); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatalf("unknown Create forgotten: %v", err) } found, cursor := false, "" for !found { - rows, err := next.ListRuntimeAllocations(t.Context(), cursor) + rows, err := deploymentStore(next).CredentialAllocations(t.Context(), cursor) if err != nil { t.Fatal(err) } @@ -85,10 +86,10 @@ func TestRuntimeAllocationAtomicOwnershipAndRecovery(t *testing.T) { if !found { t.Fatal("unknown cleanup absent from recovery scan") } - if _, err := next.SettleRuntimeCreation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, next).SettleCreation(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err := next.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, next).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } var releases int @@ -103,13 +104,13 @@ func TestRuntimeAllocationOneWinnerAndRollback(t *testing.T) { tenant, provider := uuid.NewString(), uuid.NewString() _, environment := localEnvironment(t, s, tenant) var wg sync.WaitGroup - results := make(chan RuntimeAllocation, 8) + results := make(chan deployment.Allocation, 8) errs := make(chan error, 8) for range 8 { wg.Add(1) go func() { defer wg.Done() - value, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, provider, runtimedevice.HashCredential(uuid.NewString())) + value, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, provider, runtimedevice.HashCredential(uuid.NewString())) results <- value errs <- err }() @@ -145,7 +146,7 @@ func TestRuntimeAllocationOneWinnerAndRollback(t *testing.T) { t.Cleanup(func() { _, _ = pool.Exec(context.Background(), "ALTER TABLE runtime_allocations DROP CONSTRAINT "+constraint) }) - if _, err := w.ReserveRuntimeAllocation(t.Context(), tenant, fail.ID, provider, runtimedevice.HashCredential(uuid.NewString())); err == nil { + if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: fail.ID}, provider, runtimedevice.HashCredential(uuid.NewString())); err == nil { t.Fatal("injected insert failure succeeded") } var count int @@ -159,33 +160,33 @@ func TestRuntimeAllocationExpiryAndRevocation(t *testing.T) { w := executionWriter(t, s) tenant := uuid.NewString() _, environment := localEnvironment(t, s, tenant) - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - owner, err = w.ObserveRuntimeRunning(t.Context(), owner) + owner, err = deploymentExecution(t, w).ObserveRunning(t.Context(), owner) if err != nil || !owner.CreateSettled { t.Fatalf("running observation: %+v %v", owner, err) } - if _, err := w.KeepRuntimeAllocation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).KeepAllocation(t.Context(), owner); 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 { t.Fatal(err) } - if _, err := w.KeepRuntimeAllocation(t.Context(), owner); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).KeepAllocation(t.Context(), owner); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatalf("expired allocation renewed: %v", err) } - if _, err := w.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("cleanup credential still authenticates") } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } - got, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, owner.ProviderKey, runtimedevice.HashCredential(uuid.NewString())) + got, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, owner.ProviderKey, runtimedevice.HashCredential(uuid.NewString())) if err != nil || !got.Replayed || got.State != "released" || got.KeptAt.After(time.Now()) { t.Fatalf("cleanup permitted replacement: %+v %v", got, err) } diff --git a/services/core/internal/store/runtime_capabilities_pending_test.go b/services/core/internal/store/runtime_capabilities_pending_test.go index 9ab5f5fa8..1052881c3 100644 --- a/services/core/internal/store/runtime_capabilities_pending_test.go +++ b/services/core/internal/store/runtime_capabilities_pending_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -21,6 +22,7 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { t.Fatal(err) } s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s.SetPlacement(fixtureRules(t, db)) tenant := uuid.NewString() session, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{ Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), @@ -46,7 +48,7 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { t.Fatal(err) } } - owner, err = s.GetRuntimeAllocation(t.Context(), tenant, env.ID) + owner, err = fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || initializationState(t, db.pool, owner.TenantID, owner.EnvironmentID) != "pending" || owner.State != "running" || provider.writes.Load() != 0 || provider.kills != 0 { t.Fatal("missing socket consumed initialization or requested cleanup", owner, err, provider.writes.Load(), provider.kills) } diff --git a/services/core/internal/store/runtime_compute_lifecycle_test.go b/services/core/internal/store/runtime_compute_lifecycle_test.go index 2f35bc16f..0474be944 100644 --- a/services/core/internal/store/runtime_compute_lifecycle_test.go +++ b/services/core/internal/store/runtime_compute_lifecycle_test.go @@ -15,6 +15,7 @@ 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/deployment" "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" @@ -304,7 +305,7 @@ func (f *computeLifecycleFixture) sql(query string, args ...any) { f.t.Fatal(err) } } -func (f *computeLifecycleFixture) create() (string, sessions.Session, sessions.Environment, store.RuntimeAllocation) { +func (f *computeLifecycleFixture) create() (string, sessions.Session, sessions.Environment, deployment.Allocation) { t := f.t t.Helper() tenant, session, environment := managedSession(t, f.store, f.db) @@ -321,9 +322,9 @@ func (f *computeLifecycleFixture) create() (string, sessions.Session, sessions.E owner = f.phase(tenant, environment.ID, "running") return tenant, session, environment, owner } -func (f *computeLifecycleFixture) phase(tenant, environment, phase string) store.RuntimeAllocation { +func (f *computeLifecycleFixture) phase(tenant, environment, phase string) deployment.Allocation { f.t.Helper() - var owner store.RuntimeAllocation + var owner deployment.Allocation for range 100 { ctx, cancel := context.WithTimeout(f.t.Context(), 3*time.Second) err := f.worker.ReconcileManagedRuntimes(ctx) @@ -331,7 +332,7 @@ func (f *computeLifecycleFixture) phase(tenant, environment, phase string) store if err != nil { f.t.Fatal(err) } - owner, err = f.store.GetRuntimeAllocation(f.t.Context(), tenant, environment) + owner, err = fixtureReader(f.db).EnvironmentAllocation(f.t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) if err != nil { f.t.Fatal(err) } @@ -342,13 +343,13 @@ func (f *computeLifecycleFixture) phase(tenant, environment, phase string) store f.t.Fatalf("compute phase=%s, want %s state=%s", owner.ComputePhase, phase, owner.ComputeState) return owner } -func (f *computeLifecycleFixture) complete(owner store.RuntimeAllocation) string { +func (f *computeLifecycleFixture) complete(owner deployment.Allocation) string { id := uuid.NewString() f.sql(`INSERT INTO turns(id,session_id,status,completed_at) VALUES($1,$2,'completed',clock_timestamp()-interval '2 minutes')`, id, owner.SessionID) f.sql(`UPDATE runtime_allocations SET compute_activity_at=clock_timestamp()-interval '2 minutes' WHERE id=$1`, owner.ID) return id } -func (f *computeLifecycleFixture) queued(owner store.RuntimeAllocation) string { +func (f *computeLifecycleFixture) queued(owner deployment.Allocation) string { id := uuid.NewString() f.sql(`INSERT INTO turns(id,session_id,status) VALUES($1,$2,'queued')`, id, owner.SessionID) return id @@ -439,7 +440,7 @@ func TestRuntimeComputeLifecycleQuiesceRejectionAndUnknownIntentRollback(t *test } // A restart with persisted quiescing but no acknowledgement must resume the // source. Inject only the durable phase, never a fake snapshot or new source. - current, err := f.store.GetRuntimeAllocation(t.Context(), tenant, env.ID) + current, err := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil { t.Fatal(err) } @@ -461,9 +462,10 @@ func TestRuntimeComputeLifecycleWakeDuringQuiesceResumesSource(t *testing.T) { tenant, _, env, owner := f.create() f.complete(owner) wakeResult := make(chan error, 1) + activity := fixtureDeployment(t, f.db) f.provider.mu.Lock() f.provider.beforeQuiesce = func() { - wakeResult <- f.store.TouchRuntimeActivity(t.Context(), tenant, env.ID) + wakeResult <- activity.TouchActivity(t.Context(), tenant, env.ID) } f.provider.mu.Unlock() for range 100 { @@ -499,7 +501,7 @@ func TestRuntimeComputeLifecycleSuspendedDeletionAndExpiryCleanup(t *testing.T) } else { f.sql(`UPDATE runtime_allocations SET compute_retained_until=clock_timestamp()-interval '1 second' WHERE id=$1`, owner.ID) } - reconcileManagedState(t, f.worker, f.store, tenant, env.ID, "released") + reconcileManagedState(t, f.worker, f.db, tenant, env.ID, "released") if len(f.provider.computes) != 0 || len(f.provider.snapshots) != 0 || f.provider.snapshotDeletes != 1 { t.Fatal("retained snapshot survived cleanup") } @@ -530,7 +532,7 @@ func TestRuntimeComputeLifecycleCapacityBoundsActiveAndRetained(t *testing.T) { for range 4 { f.worker.ReconcileManagedRuntimes(t.Context()) } - first, err := f.store.GetRuntimeAllocation(t.Context(), tenant, env.ID) + first, err := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || first.ComputePhase != "suspended" || f.provider.restores != 0 { t.Fatal("wake exceeded active capacity", err) } diff --git a/services/core/internal/store/runtime_configuration_cleanup_test.go b/services/core/internal/store/runtime_configuration_cleanup_test.go index d188d63bc..4911133ca 100644 --- a/services/core/internal/store/runtime_configuration_cleanup_test.go +++ b/services/core/internal/store/runtime_configuration_cleanup_test.go @@ -4,6 +4,7 @@ import ( "context" "testing" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/google/uuid" ) @@ -106,7 +107,7 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - before, err := s.GetRuntimeAllocation(t.Context(), tenant, environment.ID) + before, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || before.State == "cleanup_pending" || before.State == "released" { t.Fatal("drift authorized cleanup", before, err) } @@ -127,13 +128,13 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { if test.wantReleased { wantState = "released" } - reconcileManagedState(t, w, s, tenant, environment.ID, wantState) + reconcileManagedState(t, w, db, tenant, environment.ID, wantState) // A second observation must not turn absence after Kill into proof // that an unknown original Create can no longer mutate resources. if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - got, err := s.GetRuntimeAllocation(t.Context(), tenant, environment.ID) + got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || got.ID != owner.ID || got.State != wantState || got.CreateSettled != test.wantSettled { t.Fatal("cleanup lost ownership or settlement", got, err) } diff --git a/services/core/internal/store/runtime_connection_test.go b/services/core/internal/store/runtime_connection_test.go index ab4fa9049..088ff0eea 100644 --- a/services/core/internal/store/runtime_connection_test.go +++ b/services/core/internal/store/runtime_connection_test.go @@ -13,6 +13,7 @@ import ( "github.com/gorilla/websocket" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtime" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" @@ -106,14 +107,14 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { stop(w) w = start() assertStatus("connected") - retained, err := s.GetRuntimeAllocation(t.Context(), tenant, environment.ID) + retained, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || retained.ID != owner.ID || retained.DeviceID != owner.DeviceID || p.creates != 1 { t.Fatal("restart replaced Runtime identity", err) } if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { t.Fatal(err) } - reconcileManagedState(t, w, s, tenant, environment.ID, "released") + reconcileManagedState(t, w, db, tenant, environment.ID, "released") if conn, err := dial(p.credential); err == nil { conn.Close() t.Fatal("released Runtime reconnected") diff --git a/services/core/internal/store/runtime_creation_settlement_test.go b/services/core/internal/store/runtime_creation_settlement_test.go index 4a3560c66..008f31f04 100644 --- a/services/core/internal/store/runtime_creation_settlement_test.go +++ b/services/core/internal/store/runtime_creation_settlement_test.go @@ -5,6 +5,7 @@ import ( "errors" "testing" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/google/uuid" ) @@ -56,7 +57,7 @@ func TestManagedRuntimeConfirmedAbsentCreateReleasesAtomically(t *testing.T) { if p.kills != 0 || p.creates != 1 { t.Fatal("absence proof still called external cleanup", p.kills, p.creates) } - stored, err := s.GetRuntimeAllocation(t.Context(), tenant, environment.ID) + stored, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || stored.State != "released" || !stored.CreateSettled { t.Fatal("release not durable", err) } @@ -85,7 +86,7 @@ func TestManagedRuntimeForeignAbsenceCannotReleaseCreation(t *testing.T) { if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err == nil { t.Fatal("foreign absence accepted") } - owner, err := s.GetRuntimeAllocation(t.Context(), tenant, environment.ID) + owner, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || owner.CreateSettled || owner.State == "released" { t.Fatal("foreign proof settled original attempt", owner, err) } @@ -102,7 +103,7 @@ func TestManagedRuntimeObservedSettlementAllowsOwnedCleanup(t *testing.T) { if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { t.Fatal(err) } - reconcileManagedState(t, w, s, tenant, environment.ID, "released") + reconcileManagedState(t, w, db, tenant, environment.ID, "released") if p.kills != 1 || p.creates != 1 { t.Fatal("settled observation replayed Create or skipped cleanup", p.creates, p.kills) } diff --git a/services/core/internal/store/runtime_deployment.go b/services/core/internal/store/runtime_deployment.go deleted file mode 100644 index 080d33178..000000000 --- a/services/core/internal/store/runtime_deployment.go +++ /dev/null @@ -1,52 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -// ErrSandboxResetAdmission rejects new hosted work while a sandbox reset -// pauses admission. -var ErrSandboxResetAdmission = errors.New("hosted admission is paused for a sandbox reset") - -// sandboxProviders interprets the deployment's provider declarations for -// Session admission and placement until they move out of the Store. -var sandboxProviders = providers.Builtin() - -// New work and deployment changes share this lock. Existing receipts are checked -// first, preserving idempotent retries and cleanup while maintenance is active. -func checkRuntimeDeploymentAdmission(ctx context.Context, q *sqlc.Queries, installation string) error { - current, err := q.LockRuntimeDeployment(ctx) - if err != nil { - return err - } - if !current.InstallationID.Valid { - return nil - } - if current.ResetClear.Valid { - return ErrSandboxResetAdmission - } - if current.AdmissionPaused { - return fmt.Errorf("%w: sandbox creation is paused for provider maintenance", sessions.ErrEnvironmentUnavailable) - } - if current.ProviderKind != "" { - var spec sandbox.DeploymentSpec - if json.Unmarshal(current.Specification, &spec) != nil || sandboxProviders.ValidateSpecification(current.ProviderKind, spec) != nil { - return fmt.Errorf("%w: sandbox creation requires a deployment specification", sessions.ErrEnvironmentUnavailable) - } - } - if installation != "" { - id, err := parseConnectionGeneration(installation) - if err != nil || id != current.InstallationID { - return fmt.Errorf("%w: sandbox installation does not match deployment", sessions.ErrEnvironmentUnavailable) - } - } - return nil -} diff --git a/services/core/internal/store/runtime_deployment_test.go b/services/core/internal/store/runtime_deployment_test.go index 5beb39acb..6868bb647 100644 --- a/services/core/internal/store/runtime_deployment_test.go +++ b/services/core/internal/store/runtime_deployment_test.go @@ -11,6 +11,7 @@ import ( "github.com/google/uuid" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -73,7 +74,7 @@ func TestRuntimeDeploymentUnknownAllocationsBlockAdoptionAndSwitch(t *testing.T) old := deploymentSelection() tenant := uuid.NewString() session, environment := localEnvironment(t, s, tenant) - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } @@ -83,24 +84,24 @@ func TestRuntimeDeploymentUnknownAllocationsBlockAdoptionAndSwitch(t *testing.T) if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { t.Fatal(err) } - if _, err := w.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), owner); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("unknown creation lost cleanup ownership", err) } if err := deploymentExecution(t, w).ConfigureProcess(t.Context(), &old); err == nil { t.Fatal("deleted unknown allocation did not block adoption") } - if _, err := w.SettleRuntimeCreation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).SettleCreation(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } deploymentConfigure(t, w, &old) _, environment = localEnvironment(t, s, tenant) - owner, err = w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())) + owner, err = deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } @@ -114,11 +115,11 @@ func TestRuntimeDeploymentUnknownAllocationsBlockAdoptionAndSwitch(t *testing.T) if err := deploymentExecution(t, w).ConfigureProcess(t.Context(), nil); err == nil { t.Fatal("removing adapter orphaned unknown creation") } - replay, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())) + replay, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())) if err != nil || !replay.Replayed || replay.ID != owner.ID { t.Fatal("maintenance blocked receipt replay", replay, err) } - if _, err := w.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).RequestCleanup(t.Context(), owner); err != nil { t.Fatal("maintenance blocked cleanup", err) } } @@ -141,14 +142,14 @@ func TestRuntimeDeploymentMaintenancePreservesCreationRetriesAndOtherPlacements( t.Fatal("creation retry lost identity", err) } input.IdempotencyKey = uuid.NewString() - if _, err := s.CreateSession(t.Context(), tenant, input); !errors.Is(err, sessions.ErrEnvironmentUnavailable) { + if _, err := s.CreateSession(t.Context(), tenant, input); !errors.Is(err, placement.ErrAdmissionClosed) { t.Fatal("maintenance created hosted Session", err) } var count int if err := pool.QueryRow(t.Context(), "SELECT count(*) FROM sessions WHERE tenant_id=$1", tenant).Scan(&count); err != nil || count != 1 { t.Fatal("rejection left partial Session", count, err) } - if _, err := w.ReserveRuntimeAllocation(t.Context(), tenant, existing.Environment.ID, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, sessions.ErrEnvironmentUnavailable) { + if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: existing.Environment.ID}, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, placement.ErrAdmissionClosed) { t.Fatal("maintenance reserved new allocation", err) } for _, kind := range []string{"none", "self_hosted"} { @@ -160,10 +161,10 @@ func TestRuntimeDeploymentMaintenancePreservesCreationRetriesAndOtherPlacements( } old.AdmissionPaused = false deploymentConfigure(t, w, &old) - if _, err := w.ReserveRuntimeAllocation(t.Context(), tenant, existing.Environment.ID, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, sessions.ErrEnvironmentUnavailable) { + if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: existing.Environment.ID}, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, placement.ErrAdmissionClosed) { t.Fatal("wrong installation reserved resource", err) } - if _, err := w.ReserveRuntimeAllocation(t.Context(), tenant, existing.Environment.ID, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())); err != nil { + if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: existing.Environment.ID}, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())); err != nil { t.Fatal("resume did not reopen allocation", err) } } @@ -197,7 +198,7 @@ func TestRuntimeDeploymentMaintenanceSerializesHostedCreation(t *testing.T) { if err := tx.Commit(ctx); err != nil { t.Fatal(err) } - if err := <-done; !errors.Is(err, sessions.ErrEnvironmentUnavailable) { + if err := <-done; !errors.Is(err, placement.ErrAdmissionClosed) { t.Fatal("creation bypassed committed maintenance", err) } var count int @@ -215,7 +216,7 @@ func TestRuntimeDeploymentRetainedResourcesBlockSwitchWithoutMutation(t *testing deploymentConfigure(t, w, &old) tenant := uuid.NewString() _, environment := localEnvironment(t, s, tenant) - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, old.InstallationID, runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } @@ -274,7 +275,7 @@ func TestRuntimeDeploymentAllocationBeforeMaintenanceRetainsOwnership(t *testing } allocated := make(chan error, 1) go func() { - _, err := w.ReserveRuntimeAllocation(ctx, tenant, environment.ID, config.InstallationID, runtimedevice.HashCredential(uuid.NewString())) + _, err := deploymentExecution(t, w).ReserveAllocation(ctx, deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, config.InstallationID, runtimedevice.HashCredential(uuid.NewString())) allocated <- err }() runtimeSuspensionWaitBlocked(t, ctx, pool, blocker, allocated) @@ -292,7 +293,7 @@ func TestRuntimeDeploymentAllocationBeforeMaintenanceRetainsOwnership(t *testing if err := <-maintaining; err != nil { t.Fatal(err) } - owner, err := w.GetRuntimeAllocation(ctx, tenant, environment.ID) + owner, err := deploymentStore(w).EnvironmentAllocation(ctx, deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || owner.State != "creating" || owner.CreateSettled { t.Fatal("maintenance changed uncertain receipt", owner, err) } diff --git a/services/core/internal/store/runtime_deployment_worker_test.go b/services/core/internal/store/runtime_deployment_worker_test.go index 7306de12f..6260a2737 100644 --- a/services/core/internal/store/runtime_deployment_worker_test.go +++ b/services/core/internal/store/runtime_deployment_worker_test.go @@ -4,6 +4,7 @@ import ( "strings" "testing" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "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" @@ -45,7 +46,7 @@ func TestManagedDeploymentStartupRejectsSwitchBeforeBackendAccess(t *testing.T) if replacement.creates != 0 || replacement.kills != 0 { t.Fatal("rejected startup touched new backend") } - got, err := s.GetRuntimeAllocation(t.Context(), tenant, environment.ID) + got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || got.ID != owner.ID || got.ProviderKey != key { t.Fatal("rejected startup rewrote resource owner", got, err) } diff --git a/services/core/internal/store/runtime_environment_terminal_test.go b/services/core/internal/store/runtime_environment_terminal_test.go index e7f545584..c22bfed18 100644 --- a/services/core/internal/store/runtime_environment_terminal_test.go +++ b/services/core/internal/store/runtime_environment_terminal_test.go @@ -6,6 +6,7 @@ import ( "testing" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" @@ -25,7 +26,7 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin } reservation := initialEnvironmentReservation(t, s, pool, tenant, session.ID) writer := executionWriter(t, s) - owner, err := writer.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, writer).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } @@ -34,7 +35,7 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin t.Fatal(err) } } - if _, err := writer.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, writer).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } ended, err := s.GetSession(t.Context(), tenant, session.ID) @@ -93,7 +94,7 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin t.Fatal("expiry recorded a provisioning failure", last) } cursor, _ := s.SessionEventCursor(t.Context(), tenant, session.ID) - if _, err := writer.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, writer).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } if next, err := s.SessionEventCursor(t.Context(), tenant, session.ID); err != nil || next != cursor { @@ -102,7 +103,7 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, session.Environment.ID, uuid.NewString()); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal("late connection revived terminal environment", err) } - if _, err := writer.ReleaseRuntimeAllocation(t.Context(), owner); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, writer).ReleaseAllocation(t.Context(), owner); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("unknown creation was forgotten", err) } environmentInputHistory(t, pool, session.ID, 0, 0) @@ -120,7 +121,7 @@ func TestManagedEnvironmentFailureRollsBackWithSessionEvent(t *testing.T) { t.Fatal(err) } writer := executionWriter(t, s) - owner, err := writer.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, writer).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } @@ -131,7 +132,7 @@ func TestManagedEnvironmentFailureRollsBackWithSessionEvent(t *testing.T) { t.Cleanup(func() { _, _ = pool.Exec(context.Background(), "ALTER TABLE session_events DROP CONSTRAINT IF EXISTS "+constraint) }) - if _, err := writer.RequestRuntimeCleanup(t.Context(), owner); err == nil { + if _, err := deploymentExecution(t, writer).RequestCleanup(t.Context(), owner); err == nil { t.Fatal("cleanup committed without failure event") } current, err := s.GetSession(t.Context(), tenant, session.ID) @@ -141,7 +142,7 @@ func TestManagedEnvironmentFailureRollsBackWithSessionEvent(t *testing.T) { if reservation := initialEnvironmentReservation(t, s, pool, tenant, session.ID); reservation.State != sessions.EnvironmentInputPending { t.Fatal("partial input failure", reservation) } - allocation, err := s.GetRuntimeAllocation(t.Context(), tenant, session.Environment.ID) + allocation, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}) if err != nil || allocation.State != "creating" { t.Fatal("partial allocation transition", err) } diff --git a/services/core/internal/store/runtime_file_admission_test.go b/services/core/internal/store/runtime_file_admission_test.go index d50b8092f..5196757ee 100644 --- a/services/core/internal/store/runtime_file_admission_test.go +++ b/services/core/internal/store/runtime_file_admission_test.go @@ -9,10 +9,11 @@ import ( "github.com/google/uuid" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -func runtimeFileWriteKey(owner RuntimeAllocation) sessions.FileWriteIdentity { +func runtimeFileWriteKey(owner deployment.Allocation) sessions.FileWriteIdentity { return sessions.FileWriteIdentity{ID: uuid.NewString(), DeviceID: owner.DeviceID, RequestSHA256: strings.Repeat("a", 64)} } @@ -92,7 +93,7 @@ func TestRuntimeFileWriteAndQuiesceSerializeBothOrders(t *testing.T) { if first == "quiesce" { _, err = writes.ReserveEnvironmentFileWrite(ctx, owner.TenantID, owner.EnvironmentID, key) } else { - _, err = w.SetRuntimeCompute(ctx, owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond) + _, err = deploymentExecution(t, w).SetCompute(ctx, owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond) } done <- err }() @@ -111,10 +112,16 @@ func TestRuntimeFileWriteAndQuiesceSerializeBothOrders(t *testing.T) { if err := tx.Commit(ctx); err != nil { t.Fatal(err) } - if err := <-done; !errors.Is(err, sessions.ErrTurnConflict) { + // The late write conflicts on the Turn; the late quiesce on the + // allocation. + want := deployment.ErrAllocationConflict + if first == "quiesce" { + want = sessions.ErrTurnConflict + } + if err := <-done; !errors.Is(err, want) { t.Fatal("competing durable owner bypassed the phase fence", err) } - allocation, err := s.GetRuntimeAllocation(ctx, owner.TenantID, owner.EnvironmentID) + allocation, err := deploymentStore(s).EnvironmentAllocation(ctx, deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/runtime_idle_clock_test.go b/services/core/internal/store/runtime_idle_clock_test.go index efb1d5cbf..84e59a4b2 100644 --- a/services/core/internal/store/runtime_idle_clock_test.go +++ b/services/core/internal/store/runtime_idle_clock_test.go @@ -11,6 +11,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -18,7 +19,7 @@ import ( "github.com/jackc/pgx/v5/pgtype" ) -func managedIdleClockFixture(t *testing.T) (*Store, *Store, RuntimeAllocation) { +func managedIdleClockFixture(t *testing.T) (*Store, *Store, deployment.Allocation) { t.Helper() s, w, d := managerFixture(t, 1, 4) tenant := uuid.NewString() @@ -28,15 +29,15 @@ func managedIdleClockFixture(t *testing.T) (*Store, *Store, RuntimeAllocation) { if err != nil { t.Fatal(err) } - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, d.InstallationID, runtimedevice.HashCredential("runtime")) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, d.InstallationID, runtimedevice.HashCredential("runtime")) if err != nil { t.Fatal(err) } - owner, err = w.ObserveRuntimeRunning(t.Context(), owner) + owner, err = deploymentExecution(t, w).ObserveRunning(t.Context(), owner) if err != nil { t.Fatal(err) } - owner, err = w.SetRuntimeCompute(t.Context(), owner, "running", json.RawMessage(`{"instance":"original"}`), nil, 0) + owner, err = deploymentExecution(t, w).SetCompute(t.Context(), owner, "running", json.RawMessage(`{"instance":"original"}`), nil, 0) if err != nil { t.Fatal(err) } @@ -51,28 +52,28 @@ func runtimeDatabaseTime(t *testing.T, s *Store) time.Time { } return now } -func verifyManagedIdleClock(t *testing.T, s, w *Store, owner RuntimeAllocation, before, after time.Time) RuntimeActivity { +func verifyManagedIdleClock(t *testing.T, s, w *Store, owner deployment.Allocation, before, after time.Time) deployment.Activity { t.Helper() const idleTimeout = time.Minute // Refresh the owner so the idle policy, rather than the stale-activity fence, // must reject the recent completion. - owner, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + owner, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil { t.Fatal(err) } observedBefore := runtimeDatabaseTime(t, s) - activity, err := w.RuntimeActivity(t.Context(), owner) + activity, err := deploymentStore(w).Activity(t.Context(), owner.ID) observedAfter := runtimeDatabaseTime(t, s) if err != nil || activity.ObservedAt.Before(observedBefore) || activity.ObservedAt.After(observedAfter) || activity.ReadyToSuspend(idleTimeout) || activity.LastActivity.Before(before) || activity.LastActivity.After(after) || activity.Busy || activity.WakeRequested || !activity.HasCompletedTurn { t.Fatal("idle clock did not use committed terminal ingestion", activity, before, after, err) } until := runtimeDatabaseTime(t, s).Add(time.Hour) - if _, err := w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, idleTimeout); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, idleTimeout); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("new completion admitted premature idle", err) } // Advance only the internal activity age; the remote public timestamp remains unchanged. runtimeSuspensionSQL(t, s.pool, "UPDATE runtime_allocations SET compute_activity_at=clock_timestamp()-interval '2 minutes' WHERE id=$1", owner.ID) - if _, err := w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, idleTimeout); err != nil { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, idleTimeout); err != nil { t.Fatal("remote timestamp delayed elapsed idle timer", err) } return activity @@ -91,14 +92,14 @@ func TestManagedIdleClockIgnoresRootHostSkew(t *testing.T) { if err != nil || completed.CompletedAt.UnixMilli() != source { t.Fatal("native completion changed or rejected", completed, err) } - recorded, err := w.RuntimeActivity(t.Context(), owner) + recorded, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil { t.Fatal(err) } if _, err := w.CompleteExecution(t.Context(), owner.TenantID, owner.SessionID, turn, sessions.TurnCompleted, outcome, "", 0); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("terminal replay accepted", err) } - unchanged, err := w.RuntimeActivity(t.Context(), owner) + unchanged, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || !unchanged.LastActivity.Equal(recorded.LastActivity) { t.Fatal("terminal retry reset idle", unchanged, err) } @@ -132,14 +133,14 @@ func TestManagedIdleClockIgnoresChildHostSkewAndReplay(t *testing.T) { t.Fatal(err) } after := runtimeDatabaseTime(t, s) - recorded, err := w.RuntimeActivity(t.Context(), owner) + recorded, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil { t.Fatal(err) } if err := project(); err != nil { t.Fatal(err) } - unchanged, err := w.RuntimeActivity(t.Context(), owner) + unchanged, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || !unchanged.LastActivity.Equal(recorded.LastActivity) { t.Fatal("replayed child completion reset idle", unchanged, err) } @@ -211,12 +212,12 @@ func TestRootCompletionRejectsNonpositiveSourceTime(t *testing.T) { func TestManagedIdleClockReconnectPreservesReceipts(t *testing.T) { s, _, owner := managedIdleClockFixture(t) turn := runtimeSuspensionCompleted(t, s.pool, owner) - before, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + before, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil { t.Fatal(err) } onlineManagerNode(t, s, owner.NodeID) - after, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + after, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil || !reflect.DeepEqual(before, after) { t.Fatal("reconnect changed compute receipt or idle clock", err) } diff --git a/services/core/internal/store/runtime_initialization_test.go b/services/core/internal/store/runtime_initialization_test.go index e56fa37a1..7d61c1ebf 100644 --- a/services/core/internal/store/runtime_initialization_test.go +++ b/services/core/internal/store/runtime_initialization_test.go @@ -17,6 +17,7 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -52,6 +53,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { t.Fatal(err) } s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s.SetPlacement(fixtureRules(t, db)) tenant := uuid.NewString() input := sessions.CreateSession{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 { @@ -193,7 +195,7 @@ func TestManagedRuntimePreparationAllOperationsUsePeer(t *testing.T) { if !reflect.DeepEqual(actions, expected) || provider.commandCalls.Load() != 0 { t.Fatal("typed ordering or provider isolation", actions, provider.commandCalls.Load()) } - allocation, err := s.GetRuntimeAllocation(t.Context(), tenant, environment.ID) + allocation, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) if err != nil || initializationState(t, db.pool, allocation.TenantID, allocation.EnvironmentID) != "complete" { t.Fatal("initialization incomplete", allocation, err) } diff --git a/services/core/internal/store/runtime_input_admission_test.go b/services/core/internal/store/runtime_input_admission_test.go index afef6de5c..b80696f0a 100644 --- a/services/core/internal/store/runtime_input_admission_test.go +++ b/services/core/internal/store/runtime_input_admission_test.go @@ -7,6 +7,7 @@ import ( "github.com/google/uuid" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" @@ -23,7 +24,7 @@ func TestManagedRuntimeMaintenancePreservesCancelAndRetry(t *testing.T) { p := &lifecycleProvider{resources: map[string]sandbox.Info{}} w, stop := managedWorkerMode(t, s, db, uuid.NewString(), p, true) defer stop() - if _, err := w.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: session.Configuration}); !errors.Is(err, sessions.ErrEnvironmentUnavailable) { + if _, err := w.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: session.Configuration}); !errors.Is(err, placement.ErrAdmissionClosed) { t.Fatal("maintenance accepted new hosted Session", err) } cancel := []sessions.Input{{Kind: "cancel", Payload: json.RawMessage(`{}`)}} diff --git a/services/core/internal/store/runtime_lifecycle_nodes.go b/services/core/internal/store/runtime_lifecycle_nodes.go deleted file mode 100644 index a387dfc2b..000000000 --- a/services/core/internal/store/runtime_lifecycle_nodes.go +++ /dev/null @@ -1,126 +0,0 @@ -package store - -import ( - "context" - "errors" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -// ListRuntimeLifecycleNodes includes offline nodes: loss of connectivity never -// releases their resources. The empty identity is the single legacy lifecycle. -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 { - return nil, err - } - rows, err := s.queries.ListRuntimeLifecycleNodes(ctx) - if err != nil { - return nil, err - } - result := make([]string, 0, len(rows)) - for _, id := range rows { - result = append(result, runtimeUUID(id)) - } - return result, nil -} - -func runtimeNodePage(node, after string) (pgtype.UUID, pgtype.UUID, error) { - var nodeID pgtype.UUID - afterID := pgtype.UUID{Valid: true} - var err error - if node != "" { - nodeID, err = parseConnectionGeneration(node) - if err != nil { - return nodeID, afterID, err - } - } - if after != "" { - afterID, err = parseID(after) - } - return nodeID, afterID, err -} - -// ListRuntimeAllocationsForNode advances independently of every other node. -func (s *Store) ListRuntimeAllocationsForNode(ctx context.Context, node, after string) ([]RuntimeAllocation, error) { - nodeID, afterID, err := runtimeNodePage(node, after) - if err != nil { - return nil, err - } - ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) - defer cancel() - if err := s.checkExecutionOwnership(ctx); err != nil { - return nil, err - } - rows, err := s.queries.ListRuntimeAllocationsForNode(ctx, sqlc.ListRuntimeAllocationsForNodeParams{NodeID: nodeID, AfterID: afterID}) - if err != nil { - return nil, err - } - result := make([]RuntimeAllocation, 0, len(rows)) - for _, row := range rows { - result = append(result, runtimeAllocationFromRow(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired)) - } - return result, nil -} - -// ListUnallocatedHostedEnvironmentsForNode reads the committed placement; it -// never selects a replacement node for an unavailable original placement. -func (s *Store) ListUnallocatedHostedEnvironmentsForNode(ctx context.Context, node, after string) ([]UnallocatedHostedEnvironment, error) { - nodeID, afterID, err := runtimeNodePage(node, after) - if err != nil { - return nil, err - } - ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) - defer cancel() - if err := s.checkExecutionOwnership(ctx); err != nil { - return nil, err - } - rows, err := s.queries.ListUnallocatedHostedEnvironmentsForNode(ctx, sqlc.ListUnallocatedHostedEnvironmentsForNodeParams{NodeID: nodeID, AfterID: afterID}) - if err != nil { - return nil, err - } - result := make([]UnallocatedHostedEnvironment, 0, len(rows)) - for _, row := range rows { - result = append(result, UnallocatedHostedEnvironment{ID: runtimeUUID(row.ID), TenantID: runtimeUUID(row.TenantID)}) - } - return result, nil -} - -// ResolveRuntimeLifecycleNode routes direct provisioning before an allocation -// exists. An existing allocation must agree with its immutable placement. -func (s *Store) ResolveRuntimeLifecycleNode(ctx context.Context, tenant, environment string) (string, error) { - lookup, err := sessionpg.DeviceLookup(tenant, environment) - if err != nil { - return "", err - } - ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) - defer cancel() - if err := s.checkExecutionOwnership(ctx); err != nil { - return "", err - } - row, err := s.queries.GetRuntimeLifecyclePlacement(ctx, sqlc.GetRuntimeLifecyclePlacementParams{TenantID: lookup.TenantID, ID: lookup.ID}) - if errors.Is(err, pgx.ErrNoRows) { - return "", sessions.ErrNotFound - } - if err != nil { - return "", err - } - if row.ProviderKind == "" || row.Mode == "direct" { - if row.PlacementNodeID.Valid || row.AllocationNodeID.Valid { - return "", deployment.ErrNodeUnavailable - } - return "", nil - } - if !row.PlacementNodeID.Valid || (row.AllocationID.Valid && row.AllocationNodeID != row.PlacementNodeID) || (!row.AllocationID.Valid && row.ReleasedAt.Valid) { - return "", deployment.ErrNodeUnavailable - } - return runtimeUUID(row.PlacementNodeID), nil -} diff --git a/services/core/internal/store/runtime_lifecycle_nodes_test.go b/services/core/internal/store/runtime_lifecycle_nodes_test.go index 858b50258..248b7645f 100644 --- a/services/core/internal/store/runtime_lifecycle_nodes_test.go +++ b/services/core/internal/store/runtime_lifecycle_nodes_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" @@ -20,7 +21,7 @@ func lifecycleTestNode(t *testing.T, s *Store) string { t.Fatal(err) } id := uuid.NewString() - _, err = nodes.Enroll(t.Context(), token, deployment.Enrollment{DeploymentGeneration: 1, SpecificationDigest: SandboxDeploymentTestSpec("docker").Digest("docker"), NodeID: id, Credential: strings.Repeat("n", 64), Name: "second", Provider: "docker", BackendFingerprint: strings.Repeat("b", 64), CoreURL: s.publicURL}) + _, err = nodes.Enroll(t.Context(), token, deployment.Enrollment{DeploymentGeneration: 1, SpecificationDigest: SandboxDeploymentTestSpec("docker").Digest("docker"), NodeID: id, Credential: strings.Repeat("n", 64), Name: "second", Provider: "docker", BackendFingerprint: strings.Repeat("b", 64), CoreURL: s.placement.PublicURL()}) if err != nil { t.Fatal(err) } @@ -36,10 +37,10 @@ func lifecycleTestSession(t *testing.T, s *Store, node string) (string, sessions } return tenant, session } -func lifecycleTestAllocation(t *testing.T, s, w *Store, d deployment.ProcessDeployment, node string) RuntimeAllocation { +func lifecycleTestAllocation(t *testing.T, s, w *Store, d deployment.ProcessDeployment, node string) deployment.Allocation { t.Helper() tenant, session := lifecycleTestSession(t, s, node) - allocation, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, d.InstallationID, runtimedevice.HashCredential(uuid.NewString())) + allocation, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, d.InstallationID, runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } @@ -68,7 +69,7 @@ func TestRuntimeLifecycleNodePagesAreIndependent(t *testing.T) { var gotAlloc, gotPending []string cursor := "" for range 4 { - rows, err := w.ListRuntimeAllocationsForNode(t.Context(), node, cursor) + rows, err := deploymentStore(w).LifecycleAllocations(t.Context(), node, cursor) if err != nil { t.Fatal(err) } @@ -88,7 +89,7 @@ func TestRuntimeLifecycleNodePagesAreIndependent(t *testing.T) { } cursor = "" for range 4 { - rows, err := w.ListUnallocatedHostedEnvironmentsForNode(t.Context(), node, cursor) + rows, err := deploymentStore(w).UnallocatedEnvironments(t.Context(), node, cursor) if err != nil { t.Fatal(err) } @@ -114,11 +115,11 @@ func TestRuntimeLifecycleNodePagesAreIndependent(t *testing.T) { t.Fatal("node pagination lost or mixed rows", node, gotAlloc, gotPending) } } - if rows, err := w.ListRuntimeAllocationsForNode(t.Context(), "", ""); err != nil || len(rows) != 0 { - t.Fatal("managed allocation in legacy lane", rows, err) + if rows, err := deploymentStore(w).LifecycleAllocations(t.Context(), "", ""); err != nil || len(rows) != 0 { + t.Fatal("managed allocation in the nodeless lane", rows, err) } - if rows, err := w.ListUnallocatedHostedEnvironmentsForNode(t.Context(), "", ""); err != nil || len(rows) != 0 { - t.Fatal("managed pending in legacy lane", rows, err) + if rows, err := deploymentStore(w).UnallocatedEnvironments(t.Context(), "", ""); err != nil || len(rows) != 0 { + t.Fatal("managed pending Environment in the nodeless lane", rows, err) } } @@ -129,16 +130,16 @@ func TestRuntimeLifecycleNodeInventoryAndRouting(t *testing.T) { environment := session.Environment.ID checkRoute := func(want string, wantErr error) { t.Helper() - got, err := w.ResolveRuntimeLifecycleNode(t.Context(), tenant, environment) + got, err := deploymentService(t, w).LifecycleNode(t.Context(), tenant, environment) if got != want || !errors.Is(err, wantErr) { t.Fatal("route", got, err, want, wantErr) } } checkRoute(other, nil) // Pending has no allocation yet. - if _, err := w.ResolveRuntimeLifecycleNode(t.Context(), uuid.NewString(), environment); !errors.Is(err, sessions.ErrNotFound) { + if _, err := deploymentService(t, w).LifecycleNode(t.Context(), uuid.NewString(), environment); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("tenant boundary", err) } - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment, d.InstallationID, runtimedevice.HashCredential("runtime")) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}, d.InstallationID, runtimedevice.HashCredential("runtime")) if err != nil { t.Fatal(err) } @@ -146,14 +147,14 @@ func TestRuntimeLifecycleNodeInventoryAndRouting(t *testing.T) { t.Fatal(err) } checkRoute(other, nil) - nodes, err := w.ListRuntimeLifecycleNodes(t.Context()) + nodes, err := deploymentStore(w).LifecycleNodes(t.Context()) if err != nil || len(nodes) != 2 || !slices.Contains(nodes, other) { t.Fatal("offline node omitted", nodes, err) } if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_allocations SET node_id=NULL WHERE id=$1", owner.ID); err != nil { t.Fatal(err) } - checkRoute("", deployment.ErrNodeUnavailable) + checkRoute("", placement.ErrNodeUnavailable) if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_allocations SET node_id=$1 WHERE id=$2", other, owner.ID); err != nil { t.Fatal(err) } @@ -161,31 +162,28 @@ func TestRuntimeLifecycleNodeInventoryAndRouting(t *testing.T) { t.Fatal(err) } checkRoute(other, nil) // Deletion does not discard cleanup routing. - owner, err = w.SettleRuntimeCreation(t.Context(), owner) + owner, err = deploymentExecution(t, w).SettleCreation(t.Context(), owner) if err != nil { t.Fatal(err) } - owner, err = w.RequestRuntimeCleanup(t.Context(), owner) + owner, err = deploymentExecution(t, w).RequestCleanup(t.Context(), owner) if err != nil { t.Fatal(err) } - if _, err = w.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err = deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } checkRoute(other, nil) - if rows, err := w.ListRuntimeAllocationsForNode(t.Context(), other, ""); err != nil || len(rows) != 0 { + if rows, err := deploymentStore(w).LifecycleAllocations(t.Context(), other, ""); err != nil || len(rows) != 0 { t.Fatal("released allocation scanned", rows, err) } if err := deploymentService(t, s).RemoveNode(t.Context(), other); err != nil { t.Fatal(err) } - nodes, err = w.ListRuntimeLifecycleNodes(t.Context()) + nodes, err = deploymentStore(w).LifecycleNodes(t.Context()) if err != nil || len(nodes) != 1 || nodes[0] != d.LocalNodeID { t.Fatal("removed node discovered", nodes, err) } - if _, err := s.ListRuntimeLifecycleNodes(t.Context()); err == nil { - t.Fatal("unleased inventory accepted") - } } func TestRuntimeLifecycleNodeRejectsMissingOrReleasedPlacement(t *testing.T) { @@ -196,10 +194,10 @@ func TestRuntimeLifecycleNodeRejectsMissingOrReleasedPlacement(t *testing.T) { if _, err := s.pool.Exec(t.Context(), mutation, session.Environment.ID); err != nil { t.Fatal(err) } - if _, err := w.ResolveRuntimeLifecycleNode(t.Context(), tenant, session.Environment.ID); !errors.Is(err, deployment.ErrNodeUnavailable) { + if _, err := deploymentService(t, w).LifecycleNode(t.Context(), tenant, session.Environment.ID); !errors.Is(err, placement.ErrNodeUnavailable) { t.Fatal("invalid placement routed", err) } - rows, err := w.ListUnallocatedHostedEnvironmentsForNode(t.Context(), d.LocalNodeID, "") + rows, err := deploymentStore(w).UnallocatedEnvironments(t.Context(), d.LocalNodeID, "") if err != nil || len(rows) != 0 { t.Fatal("invalid placement provisioned", rows, err) } @@ -207,45 +205,30 @@ func TestRuntimeLifecycleNodeRejectsMissingOrReleasedPlacement(t *testing.T) { } } -func TestRuntimeLifecycleLegacyLaneAndOwnerLoss(t *testing.T) { - s, w, _, a := legacyAdoptionFixture(t) - nodes, err := w.ListRuntimeLifecycleNodes(t.Context()) +func TestRuntimeLifecycleNodelessLane(t *testing.T) { + s, w, _, a := nodelessAllocationFixture(t) + nodes, err := deploymentStore(w).LifecycleNodes(t.Context()) if err != nil || !slices.Equal(nodes, []string{""}) { t.Fatal(nodes, err) } - node, err := w.ResolveRuntimeLifecycleNode(t.Context(), a.TenantID, a.EnvironmentID) + node, err := deploymentService(t, w).LifecycleNode(t.Context(), a.TenantID, a.EnvironmentID) if err != nil || node != "" { t.Fatal(node, err) } - rows, err := w.ListRuntimeAllocationsForNode(t.Context(), "", "") + rows, err := deploymentStore(w).LifecycleAllocations(t.Context(), "", "") if err != nil || len(rows) != 1 || rows[0].ID != a.ID { t.Fatal(rows, err) } tenant := uuid.NewString() _, e := localEnvironment(t, s, tenant) - pending, err := w.ListUnallocatedHostedEnvironmentsForNode(t.Context(), "", "") + pending, err := deploymentStore(w).UnallocatedEnvironments(t.Context(), "", "") if err != nil || len(pending) != 1 || pending[0].ID != e.ID { t.Fatal(pending, err) } - if _, err := w.ListRuntimeAllocationsForNode(t.Context(), "bad", ""); err == nil { + if _, err := deploymentStore(w).LifecycleAllocations(t.Context(), "bad", ""); err == nil { t.Fatal("invalid node accepted") } - if _, err := w.ListUnallocatedHostedEnvironmentsForNode(t.Context(), "", "bad"); err == nil { + if _, err := deploymentStore(w).UnallocatedEnvironments(t.Context(), "", "bad"); err == nil { t.Fatal("invalid cursor accepted") } - if err := w.lease.Close(t.Context()); err != nil { - t.Fatal(err) - } - if _, err := w.ListRuntimeLifecycleNodes(t.Context()); err == nil { - t.Fatal("closed owner enumerated") - } - if _, err := w.ListRuntimeAllocationsForNode(t.Context(), "", ""); err == nil { - t.Fatal("closed owner scanned") - } - if _, err := w.ListUnallocatedHostedEnvironmentsForNode(t.Context(), "", ""); err == nil { - t.Fatal("closed owner provisioned") - } - if _, err := w.ResolveRuntimeLifecycleNode(t.Context(), a.TenantID, a.EnvironmentID); err == nil { - t.Fatal("closed owner routed") - } } diff --git a/services/core/internal/store/runtime_lifecycle_test.go b/services/core/internal/store/runtime_lifecycle_test.go index c7cd5e3eb..1473b627a 100644 --- a/services/core/internal/store/runtime_lifecycle_test.go +++ b/services/core/internal/store/runtime_lifecycle_test.go @@ -12,6 +12,7 @@ import ( "github.com/google/uuid" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "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" @@ -127,13 +128,13 @@ func managedSession(t *testing.T, s *store.Store, db fixtureDB) (string, session return tenant, v, env } -func reconcileManagedState(t *testing.T, w *execution.Worker, s *store.Store, tenant, environment, state string) { +func reconcileManagedState(t *testing.T, w *execution.Worker, db fixtureDB, tenant, environment, state string) { t.Helper() for range 100 { if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - got, err := s.GetRuntimeAllocation(t.Context(), tenant, environment) + got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) if err != nil { t.Fatal(err) } @@ -160,8 +161,8 @@ func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { } stop() next, _ := managedWorker(t, s, db, key, p) - reconcileManagedState(t, next, s, tenant, env.ID, "running") - recovered, err := s.GetRuntimeAllocation(t.Context(), tenant, env.ID) + reconcileManagedState(t, next, db, tenant, env.ID, "running") + recovered, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || recovered.ID != owner.ID || !recovered.CreateSettled || recovered.State != "running" { t.Fatalf("lost response recovery: %+v %v", recovered, err) } @@ -173,8 +174,8 @@ func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { t.Fatal(err) } // A scan may first exhaust its previous cursor before starting a new cycle. - reconcileManagedState(t, next, s, tenant, env.ID, "released") - clean, err := s.GetRuntimeAllocation(t.Context(), tenant, env.ID) + reconcileManagedState(t, next, db, tenant, env.ID, "released") + clean, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || clean.State != "released" || p.kills != 1 { t.Fatalf("deleted cleanup: %+v %v", clean, err) } @@ -196,8 +197,8 @@ func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { t.Fatal(err) } - reconcileManagedState(t, w, s, tenant, env.ID, "cleanup_pending") - got, err := s.GetRuntimeAllocation(t.Context(), tenant, env.ID) + reconcileManagedState(t, w, db, tenant, env.ID, "cleanup_pending") + got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "cleanup_pending" || got.CreateSettled || p.creates != 1 { t.Fatalf("unknown creation forgotten: %+v %v", got, err) } @@ -206,8 +207,8 @@ func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { } // A late completion is still owned and reclaimed on the next scan. p.resources[owner.ID] = sandbox.Info{Reference: sandbox.Reference{TenantID: tenant, EnvironmentID: env.ID, AllocationID: owner.ID}, ProviderID: owner.ID, State: "running", BootstrapComplete: true} - reconcileManagedState(t, w, s, tenant, env.ID, "released") - got, err = s.GetRuntimeAllocation(t.Context(), tenant, env.ID) + reconcileManagedState(t, w, db, tenant, env.ID, "released") + got, err = fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "released" || len(p.resources) != 0 { t.Fatalf("late creation escaped cleanup: %+v %v", got, err) } @@ -227,8 +228,8 @@ func TestManagedRuntimeExpiryRevokesWhenProviderUnavailable(t *testing.T) { t.Fatal(err) } p.unavailable = true - reconcileManagedState(t, w, s, tenant, env.ID, "cleanup_pending") - got, err := s.GetRuntimeAllocation(t.Context(), tenant, env.ID) + reconcileManagedState(t, w, db, tenant, env.ID, "cleanup_pending") + got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "cleanup_pending" { t.Fatalf("expiry lost on provider failure: %+v %v", got, err) } @@ -266,7 +267,7 @@ func TestManagedRuntimeStoppedComputeDoesNotRequestCleanup(t *testing.T) { if p.gets == before { t.Fatal("fixture allocation not inspected") } - got, err := s.GetRuntimeAllocation(t.Context(), tenant, env.ID) + got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "running" || p.kills != 0 || p.creates != 1 { t.Fatalf("compute interruption authorized replacement/cleanup: %+v %v", got, err) } diff --git a/services/core/internal/store/runtime_node_generations_test.go b/services/core/internal/store/runtime_node_generations_test.go index 0ea57bc7c..aa7e857c7 100644 --- a/services/core/internal/store/runtime_node_generations_test.go +++ b/services/core/internal/store/runtime_node_generations_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -113,14 +114,14 @@ func TestNodeGenerationsCapacityFallbackAndImmutablePending(t *testing.T) { t.Fatal("late readiness moved pin or erased serving readiness", n) } } - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, pending.Environment.ID, first.InstallationID, runtimedevice.HashCredential("runtime")) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: pending.Environment.ID}, first.InstallationID, runtimedevice.HashCredential("runtime")) if err != nil { t.Fatal(err) } if owner.NodeID != pendingNode || owner.DeploymentGeneration != uint64(pendingGeneration) { t.Fatal("pending placement moved", owner) } - n, g, err := s.ResolveRuntimeGeneration(t.Context(), sandbox.Reference{TenantID: tenant, EnvironmentID: owner.EnvironmentID, AllocationID: owner.ID}) + n, g, err := deploymentService(t, s).AllocationGeneration(t.Context(), sandbox.Reference{TenantID: tenant, EnvironmentID: owner.EnvironmentID, AllocationID: owner.ID}) if err != nil || n != pendingNode || g != 1 { t.Fatal(n, g, err) } @@ -162,7 +163,7 @@ func TestNodeGenerationsReconnectAndV1Fallback(t *testing.T) { if err != nil || nodes[0].ProviderReady || *nodes[0].Rollout.ReadyGeneration != 1 { t.Fatal("reconnect inherited readiness or lost pin", nodes, err) } - if _, err = s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, deployment.ErrNodeUnavailable) { + if _, err = s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, placement.ErrNodeUnavailable) { t.Fatal("unconfirmed connection admitted", err) } if err = service.Heartbeat(t.Context(), node.NodeID, connection, first.OwnerEpoch, deployment.NodeHealth{ProviderReady: true}); err != nil { @@ -182,7 +183,7 @@ func TestNodeGenerationPreparationRefusalCreatesNoProvisionalOwnership(t *testin } generationHeartbeat(t, s, node, connection, first, "preparing") tenant := uuid.NewString() - if _, err := s.CreateSession(t.Context(), tenant, managerSessionInput(uuid.NewString())); !errors.Is(err, deployment.ErrNodesPreparing) { + if _, err := s.CreateSession(t.Context(), tenant, managerSessionInput(uuid.NewString())); !errors.Is(err, placement.ErrNodesPreparing) { t.Fatal("actual preparation was not identified", err) } var sessions, placements int @@ -190,7 +191,7 @@ func TestNodeGenerationPreparationRefusalCreatesNoProvisionalOwnership(t *testin t.Fatal("refusal left provisional ownership", sessions, placements, err) } generationHeartbeat(t, s, node, connection, first, "failed") - if _, err := s.CreateSession(t.Context(), tenant, managerSessionInput(uuid.NewString())); !errors.Is(err, deployment.ErrNodeUnavailable) { + if _, err := s.CreateSession(t.Context(), tenant, managerSessionInput(uuid.NewString())); !errors.Is(err, placement.ErrNodeUnavailable) { t.Fatal("failed preparation advertised active work", 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 23c4776c6..1a7f7fb14 100644 --- a/services/core/internal/store/runtime_node_lifecycle_fixture_test.go +++ b/services/core/internal/store/runtime_node_lifecycle_fixture_test.go @@ -78,6 +78,7 @@ func (p *nodeIsolationProvider) RunCommand(ctx context.Context, r sandbox.Refere type nodeIsolationFixture struct { t *testing.T store *store.Store + db fixtureDB nodes *deployment.Service pool *pgxpool.Pool worker *execution.Worker @@ -98,6 +99,7 @@ func newNodeIsolationFixture(t *testing.T, mode string) *nodeIsolationFixture { t.Fatal(err) } s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s.SetPlacement(fixtureRules(t, db)) 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{})} @@ -125,7 +127,7 @@ func newNodeIsolationFixture(t *testing.T, mode string) *nodeIsolationFixture { cp.mu.Unlock() server.Close() }) - f := &nodeIsolationFixture{initializationCancel: cancelPreparation, t: t, store: s, nodes: fixtureDeployment(t, db), pool: pool, provider: p, key: uuid.NewString(), nodeA: uuid.NewString(), nodeB: uuid.NewString()} + f := &nodeIsolationFixture{initializationCancel: cancelPreparation, t: t, store: s, db: db, nodes: fixtureDeployment(t, db), pool: pool, provider: p, key: uuid.NewString(), nodeA: uuid.NewString(), nodeB: uuid.NewString()} // Keep restored compute awake throughout the isolation assertions. // The suspension setup explicitly dates its activity two minutes in the past. policy := &execution.RuntimeSuspensionPolicy{IdleTimeout: time.Minute, Retention: time.Hour, MaxActive: 100, MaxRetained: 100} @@ -220,7 +222,7 @@ func (f *nodeIsolationFixture) session(node string, initialize bool) (string, se } return tenant, session, env } -func (f *nodeIsolationFixture) provision(tenant string, env sessions.Environment) store.RuntimeAllocation { +func (f *nodeIsolationFixture) provision(tenant string, env sessions.Environment) deployment.Allocation { f.t.Helper() ctx, cancel := context.WithTimeout(f.t.Context(), 3*time.Second) defer cancel() @@ -230,13 +232,13 @@ func (f *nodeIsolationFixture) provision(tenant string, env sessions.Environment } return owner } -func (f *nodeIsolationFixture) phase(tenant, environment, phase string) store.RuntimeAllocation { +func (f *nodeIsolationFixture) phase(tenant, environment, phase string) deployment.Allocation { f.t.Helper() for range 10 { if err := f.worker.ReconcileManagedRuntimes(f.t.Context()); err != nil { f.t.Fatal(err) } - owner, err := f.store.GetRuntimeAllocation(f.t.Context(), tenant, environment) + owner, err := fixtureReader(f.db).EnvironmentAllocation(f.t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) if err != nil { f.t.Fatal(err) } @@ -245,7 +247,7 @@ func (f *nodeIsolationFixture) phase(tenant, environment, phase string) store.Ru } } f.t.Fatal("manual lifecycle did not reach " + phase) - return store.RuntimeAllocation{} + return deployment.Allocation{} } func (f *nodeIsolationFixture) run() { f.t.Helper() diff --git a/services/core/internal/store/runtime_node_lifecycle_test.go b/services/core/internal/store/runtime_node_lifecycle_test.go index 8046c574a..7055fe175 100644 --- a/services/core/internal/store/runtime_node_lifecycle_test.go +++ b/services/core/internal/store/runtime_node_lifecycle_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/deployment" "github.com/google/uuid" ) @@ -47,7 +47,7 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { case <-time.After(8 * time.Second): t.Fatalf("node A did not enter blocked %s operation", mode) } - online, err := f.store.RuntimeNodeAvailable(t.Context(), f.nodeA) + online, err := fixtureReader(f.db).NodeOnline(t.Context(), f.nodeA) if err != nil || !online { t.Fatal("test node must stay online while its helper is blocked", err) } @@ -61,7 +61,7 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { f.provider.mu.Unlock() tenant, _, env := f.session(f.nodeB, true) var callers sync.WaitGroup - results := make(chan store.RuntimeAllocation, 12) + results := make(chan deployment.Allocation, 12) failures := make(chan error, 12) for range 12 { callers.Add(1) @@ -97,7 +97,7 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { if created != 1 { t.Fatalf("concurrent direct/scan replayed Create: %d", created) } - if err := f.store.TouchRuntimeActivity(t.Context(), wakeTenant, wakeEnv.ID); err != nil { + if err := f.nodes.TouchActivity(t.Context(), wakeTenant, wakeEnv.ID); err != nil { t.Fatal(err) } if err := f.store.DeleteSession(t.Context(), deleteTenant, deleteSession.ID); err != nil { @@ -107,9 +107,9 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { // Only independent normal five-second loops drive these transitions. // No manual reconciliation or wake hint accelerates healthy nodes. waitNodeIsolation(t, 18*time.Second, func() (bool, string) { - wake, e1 := f.store.GetRuntimeAllocation(t.Context(), wakeTenant, wakeEnv.ID) - deleted, e2 := f.store.GetRuntimeAllocation(t.Context(), deleteTenant, deleteEnv.ID) - initialized, e3 := f.store.GetRuntimeAllocation(t.Context(), tenant, env.ID) + wake, e1 := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: wakeTenant, EnvironmentID: wakeEnv.ID}) + deleted, e2 := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: deleteTenant, EnvironmentID: deleteEnv.ID}) + initialized, e3 := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) f.provider.mu.Lock() restores := f.provider.restores f.provider.mu.Unlock() @@ -139,7 +139,7 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { t.Fatal(err) } waitNodeIsolation(t, 7*time.Second, func() (bool, string) { - owner, err := f.store.GetRuntimeAllocation(t.Context(), ct, ce.ID) + owner, err := fixtureReader(f.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: ct, EnvironmentID: ce.ID}) return err == nil && owner.State == "released", fmt.Sprintf("new node allocation=%s/%s err=%v", owner.State, owner.ComputePhase, err) }) if err := f.nodes.RemoveNode(t.Context(), nodeC); err != nil { diff --git a/services/core/internal/store/runtime_nodes.go b/services/core/internal/store/runtime_nodes.go index 6e42de4b1..ef13d1565 100644 --- a/services/core/internal/store/runtime_nodes.go +++ b/services/core/internal/store/runtime_nodes.go @@ -1,62 +1,13 @@ package store import ( - "context" - "errors" - "time" - "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -type RuntimeNodeAllocation struct { - DeploymentGeneration uint64 `json:"deployment_generation"` - Diagnostic string `json:"diagnostic"` - ID string `json:"id"` - NodeID string `json:"node_id"` - TenantID string `json:"tenant_id"` - SessionID string `json:"session_id"` - EnvironmentID string `json:"environment_id"` - State string `json:"state"` - ComputePhase string `json:"compute_phase"` - // The time the allocation entered its current compute_phase, or null when unknown; an allocation that existed before Core recorded it reports null until its next phase change. For a suspended microsandbox allocation, this time plus the deployment's snapshot retention tells roughly when Core reclaims it. - ComputePhaseChangedAt *time.Time `json:"compute_phase_changed_at" extensions:"x-nullable"` - Initialization string `json:"initialization"` - CreatedAt time.Time `json:"created_at"` -} - func runtimeUUID(id pgtype.UUID) string { if !id.Valid { return "" } return uuid.UUID(id.Bytes).String() } - -func (s *Store) ListNodeRuntimeAllocations(ctx context.Context, nodeID string) ([]RuntimeNodeAllocation, error) { - id, err := parseConnectionGeneration(nodeID) - if err != nil { - return nil, err - } - if _, err := s.queries.GetRuntimeNode(ctx, id); errors.Is(err, pgx.ErrNoRows) { - return nil, sessions.ErrNotFound - } else if err != nil { - return nil, err - } - rows, err := s.queries.ListNodeRuntimeAllocations(ctx, id) - if err != nil { - return nil, err - } - out := make([]RuntimeNodeAllocation, 0, len(rows)) - for _, a := range rows { - var phaseChanged *time.Time - if a.ComputePhaseChangedAt.Valid { - value := a.ComputePhaseChangedAt.Time - phaseChanged = &value - } - out = append(out, RuntimeNodeAllocation{DeploymentGeneration: uint64(a.DeploymentGeneration.Int64), Diagnostic: a.ObservationError, ID: runtimeUUID(a.ID), NodeID: runtimeUUID(a.NodeID), TenantID: runtimeUUID(a.TenantID), SessionID: runtimeUUID(a.SessionID), EnvironmentID: runtimeUUID(a.EnvironmentID), State: a.State, ComputePhase: a.ComputePhase, ComputePhaseChangedAt: phaseChanged, Initialization: a.Initialization, CreatedAt: a.CreatedAt.Time}) - } - return out, nil -} diff --git a/services/core/internal/store/runtime_nodes_test.go b/services/core/internal/store/runtime_nodes_test.go index e075cdbb3..7627a1421 100644 --- a/services/core/internal/store/runtime_nodes_test.go +++ b/services/core/internal/store/runtime_nodes_test.go @@ -10,6 +10,7 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" @@ -135,7 +136,7 @@ func TestRuntimeNodesAtomicPlacementAndRetry(t *testing.T) { for err := range failures { if err == nil { successes++ - } else if !errors.Is(err, deployment.ErrNodeUnavailable) { + } else if !errors.Is(err, placement.ErrNodeUnavailable) { t.Fatal(err) } } @@ -186,7 +187,7 @@ func TestRuntimeNodesEnrollmentAndEpoch(t *testing.T) { if err != nil { t.Fatal(err) } - input := deployment.Enrollment{DeploymentGeneration: 1, SpecificationDigest: SandboxDeploymentTestSpec("docker").Digest("docker"), NodeID: uuid.NewString(), Credential: strings.Repeat("x", 64), Name: "remote", Provider: "microsandbox", BackendFingerprint: strings.Repeat("b", 64), CoreURL: s.publicURL} + input := deployment.Enrollment{DeploymentGeneration: 1, SpecificationDigest: SandboxDeploymentTestSpec("docker").Digest("docker"), NodeID: uuid.NewString(), Credential: strings.Repeat("x", 64), Name: "remote", Provider: "microsandbox", BackendFingerprint: strings.Repeat("b", 64), CoreURL: s.placement.PublicURL()} if _, err := nodes.Enroll(t.Context(), token, input); !errors.Is(err, deployment.ErrInvalidInput) { t.Fatal("mixed provider accepted", err) } @@ -222,7 +223,7 @@ func TestRuntimeNodesEnrollmentAndEpoch(t *testing.T) { if err := nodes.Heartbeat(t.Context(), input.NodeID, connection, epoch, deployment.NodeHealth{ProviderReady: true}); !errors.Is(err, deployment.ErrNodeCredential) { t.Fatal("old epoch heartbeat revived node", err) } - if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput("stale")); !errors.Is(err, deployment.ErrNodeUnavailable) { + if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput("stale")); !errors.Is(err, placement.ErrNodeUnavailable) { t.Fatal("stale node admitted", err) } onlineManagerNode(t, s, d.LocalNodeID) @@ -241,7 +242,7 @@ func TestRuntimeNodesRetention(t *testing.T) { if err != nil { t.Fatal(err) } - retained, err := w.ReserveRuntimeAllocation(t.Context(), tenant, first.Environment.ID, next.InstallationID, runtimedevice.HashCredential("runtime")) + retained, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: first.Environment.ID}, next.InstallationID, runtimedevice.HashCredential("runtime")) if err != nil { t.Fatal(err) } @@ -258,18 +259,18 @@ func TestRuntimeNodesRetention(t *testing.T) { if err := s.DeleteSession(t.Context(), tenant, first.ID); err != nil { t.Fatal(err) } - retained, err = w.RequestRuntimeCleanup(t.Context(), retained) + retained, err = deploymentExecution(t, w).RequestCleanup(t.Context(), retained) if err != nil { t.Fatal(err) } if err := nodes.RemoveNode(t.Context(), next.LocalNodeID); !errors.Is(err, deployment.ErrNodeInUse) { t.Fatal("unknown cleanup released node", err) } - retained, err = w.SettleRuntimeCreation(t.Context(), retained) + retained, err = deploymentExecution(t, w).SettleCreation(t.Context(), retained) if err != nil { t.Fatal(err) } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), retained); err != nil { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), retained); err != nil { t.Fatal(err) } if err := nodes.RemoveNode(t.Context(), next.LocalNodeID); !errors.Is(err, deployment.ErrLocalNodeConfigured) { @@ -297,14 +298,14 @@ func TestRuntimeNodesRestoreAndCreationShareCapacity(t *testing.T) { if err != nil { t.Fatal(err) } - allocation, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, d.InstallationID, runtimedevice.HashCredential("runtime")) + allocation, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, d.InstallationID, runtimedevice.HashCredential("runtime")) if err != nil { t.Fatal(err) } if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_allocations SET state='running',create_settled=true,compute_phase='suspended',compute_retained_until=clock_timestamp()+interval '1 hour',compute_state=$2::jsonb WHERE id=$1", allocation.ID, json.RawMessage(`{"snapshot":{"id":"owned"}}`)); err != nil { t.Fatal(err) } - allocation, err = s.GetRuntimeAllocation(t.Context(), tenant, session.Environment.ID) + allocation, err = deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}) if err != nil { t.Fatal(err) } @@ -313,7 +314,7 @@ func TestRuntimeNodesRestoreAndCreationShareCapacity(t *testing.T) { results := make(chan error, 2) go func() { <-start - _, err := w.SetRuntimeCompute(t.Context(), allocation, "restoring", json.RawMessage(`{"target":{"id":"restore"}}`), &until, 0) + _, err := deploymentExecution(t, w).SetCompute(t.Context(), allocation, "restoring", json.RawMessage(`{"target":{"id":"restore"}}`), &until, 0) results <- err }() go func() { @@ -327,7 +328,7 @@ func TestRuntimeNodesRestoreAndCreationShareCapacity(t *testing.T) { err := <-results if err == nil { success++ - } else if !errors.Is(err, deployment.ErrNodeUnavailable) { + } else if !errors.Is(err, placement.ErrNodeUnavailable) { t.Fatal(err) } } @@ -355,11 +356,11 @@ func TestRuntimeNodesLongOfflineRetainsExactAllocation(t *testing.T) { if err != nil { t.Fatal(err) } - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, d.InstallationID, runtimedevice.HashCredential("runtime")) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, d.InstallationID, runtimedevice.HashCredential("runtime")) if err != nil { t.Fatal(err) } - owner, err = w.ObserveRuntimeRunning(t.Context(), owner) + owner, err = deploymentExecution(t, w).ObserveRunning(t.Context(), owner) if err != nil { t.Fatal(err) } @@ -369,7 +370,7 @@ func TestRuntimeNodesLongOfflineRetainsExactAllocation(t *testing.T) { if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_nodes SET connection_id=NULL WHERE id=$1", d.LocalNodeID); err != nil { t.Fatal(err) } - offline, err := s.GetRuntimeAllocation(t.Context(), tenant, session.Environment.ID) + offline, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}) if err != nil || offline.Expired || offline.State != "running" { t.Fatal("offline treated as destructive expiry", offline, err) } @@ -382,11 +383,11 @@ func TestRuntimeNodesLongOfflineRetainsExactAllocation(t *testing.T) { t.Fatal("lost local state created replacement identity") } onlineManagerNode(t, s, d.LocalNodeID) - resumed, err := w.ObserveRuntimeRunning(t.Context(), offline) + resumed, err := deploymentExecution(t, w).ObserveRunning(t.Context(), offline) if err != nil || resumed.ID != owner.ID || resumed.DeviceID != owner.DeviceID || resumed.NodeID != owner.NodeID { t.Fatal("reconnect changed instance", resumed, err) } - if _, err := w.KeepRuntimeAllocation(t.Context(), resumed); err != nil { + if _, err := deploymentExecution(t, w).KeepAllocation(t.Context(), resumed); err != nil { t.Fatal("offline observation lease could not renew", err) } if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_allocations SET compute_phase='suspended',compute_state=$2::jsonb,compute_retained_until=clock_timestamp()+interval '1 day' WHERE id=$1", owner.ID, json.RawMessage(`{"snapshot":{"id":"same-snapshot"}}`)); err != nil { @@ -395,12 +396,12 @@ func TestRuntimeNodesLongOfflineRetainsExactAllocation(t *testing.T) { if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_nodes SET connection_id=NULL WHERE id=$1", d.LocalNodeID); err != nil { t.Fatal(err) } - retained, err := s.GetRuntimeAllocation(t.Context(), tenant, session.Environment.ID) + retained, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}) if err != nil || retained.Expired || string(retained.ComputeState) != `{"snapshot": {"id": "same-snapshot"}}` { t.Fatal(retained, err) } onlineManagerNode(t, s, d.LocalNodeID) - same, err := s.GetRuntimeAllocation(t.Context(), tenant, session.Environment.ID) + same, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}) if err != nil || same.ID != owner.ID || string(same.ComputeState) != string(retained.ComputeState) { t.Fatal("snapshot changed across reconnect", same, err) } @@ -411,7 +412,7 @@ func TestRuntimeNodesLongOfflineRetainsExactAllocation(t *testing.T) { if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_allocations SET compute_retained_until=clock_timestamp()-interval '1 second' WHERE id=$1", owner.ID); err != nil { t.Fatal(err) } - expired, err := s.GetRuntimeAllocation(t.Context(), tenant, session.Environment.ID) + expired, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}) if err != nil || !expired.Expired { t.Fatal("explicit snapshot retention ignored", expired, err) } diff --git a/services/core/internal/store/runtime_observation.go b/services/core/internal/store/runtime_observation.go deleted file mode 100644 index 6c0db28ed..000000000 --- a/services/core/internal/store/runtime_observation.go +++ /dev/null @@ -1,49 +0,0 @@ -package store - -import ( - "context" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5/pgtype" -) - -// RecordRuntimeObservation changes diagnostics only for the observed lifecycle revision. -// It never releases ownership, changes public readiness or authorizes replacement. -func (s *Store) RecordRuntimeObservation(ctx context.Context, owner RuntimeAllocation, diagnostic string) error { - switch diagnostic { - case "", "node_unavailable", "resource_missing", "compute_unconfirmed", "ownership_mismatch", "provider_unavailable": - default: - return sessions.ErrInvalidInput - } - if owner.NodeID == "" { - return nil - } - _, err := s.mutateRuntimeAllocation(ctx, owner, false, func(ctx context.Context, q *sqlc.Queries, row sqlc.RuntimeAllocation) (sqlc.RuntimeAllocation, error) { - if row.ComputeRevision != owner.ComputeRevision || row.State != owner.State || row.State == "released" { - return row, nil - } - err := q.SetRuntimeObservation(ctx, sqlc.SetRuntimeObservationParams{ID: row.ID, ComputeRevision: owner.ComputeRevision, State: owner.State, ObservationError: diagnostic}) - return row, err - }) - return err -} -func (s *Store) RuntimeNodeAvailable(ctx context.Context, node string) (bool, error) { - id, err := parseConnectionGeneration(node) - if err != nil { - return false, err - } - return runtimeNodeAvailable(ctx, s.queries, id) -} -func runtimeNodeAvailable(ctx context.Context, q *sqlc.Queries, node pgtype.UUID) (bool, error) { - nodes, err := q.ListRuntimeNodes(ctx, pgtype.UUID{}) - if err != nil { - return false, err - } - for _, n := range nodes { - if n.ID == node { - return n.Online, nil - } - } - return false, nil -} diff --git a/services/core/internal/store/runtime_observation_scan_test.go b/services/core/internal/store/runtime_observation_scan_test.go index 4aca54d61..9bc3238b3 100644 --- a/services/core/internal/store/runtime_observation_scan_test.go +++ b/services/core/internal/store/runtime_observation_scan_test.go @@ -3,8 +3,6 @@ package store_test import ( "slices" "testing" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *testing.T) { @@ -15,7 +13,7 @@ func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *test expected = append(expected, session.ID) } slices.Sort(expected) - if err := s.DeleteSession(t.Context(), sessionTenant(t, s, expected[2]), expected[2]); err != nil { + if err := s.DeleteSession(t.Context(), sessionTenant(t, db, expected[2]), expected[2]); err != nil { t.Fatal(err) } expected = append(expected[:2], expected[3:]...) @@ -23,7 +21,7 @@ func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *test var got []string cursor := "" for { - page, err := s.ListRuntimeObservationSessions(t.Context(), cursor, 2) + page, err := fixtureReader(db).ObservationSessions(t.Context(), cursor, 2) if err != nil { t.Fatal(err) } @@ -47,16 +45,16 @@ func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *test if !slices.Equal(got, expected) { t.Fatalf("observation scan = %v, want %v", got, expected) } - if _, err := s.ListRuntimeObservationSessions(t.Context(), "", 0); err == nil { + if _, err := fixtureReader(db).ObservationSessions(t.Context(), "", 0); err == nil { t.Fatal("zero observation page size was accepted") } } -func sessionTenant(t *testing.T, s *store.Store, sessionID string) string { +func sessionTenant(t *testing.T, db fixtureDB, sessionID string) string { t.Helper() // The deployment-wide scan intentionally discovers tenant identity without // enumerating configured API keys. Use that same read to locate this fixture. - page, err := s.ListRuntimeObservationSessions(t.Context(), "", 100) + page, err := fixtureReader(db).ObservationSessions(t.Context(), "", 100) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/runtime_observation_test.go b/services/core/internal/store/runtime_observation_test.go index b072b1464..84a2fa366 100644 --- a/services/core/internal/store/runtime_observation_test.go +++ b/services/core/internal/store/runtime_observation_test.go @@ -6,7 +6,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" ) @@ -17,18 +16,18 @@ func TestRuntimeNodeObservationRetainsResourcesAndFencesStaleResults(t *testing. if err != nil { t.Fatal(err) } - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, d.InstallationID, runtimedevice.HashCredential("runtime")) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, d.InstallationID, runtimedevice.HashCredential("runtime")) if err != nil { t.Fatal(err) } - owner, err = w.ObserveRuntimeRunning(t.Context(), owner) + owner, err = deploymentExecution(t, w).ObserveRunning(t.Context(), owner) if err != nil { t.Fatal(err) } - if err := w.RecordRuntimeObservation(t.Context(), owner, "node_unavailable"); err != nil { + if err := deploymentExecution(t, w).RecordObservation(t.Context(), owner, "node_unavailable"); err != nil { t.Fatal(err) } - retained, err := s.GetRuntimeAllocation(t.Context(), tenant, session.Environment.ID) + retained, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}) if err != nil || retained.State != "running" || retained.ID != owner.ID || retained.ObservationError != "node_unavailable" { t.Fatal(retained, err) } @@ -39,24 +38,24 @@ func TestRuntimeNodeObservationRetainsResourcesAndFencesStaleResults(t *testing. if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_allocations SET compute_revision=compute_revision+1,observation_error='resource_missing' WHERE id=$1", owner.ID); err != nil { t.Fatal(err) } - if err := w.RecordRuntimeObservation(t.Context(), owner, ""); err != nil { + if err := deploymentExecution(t, w).RecordObservation(t.Context(), owner, ""); err != nil { t.Fatal(err) } - current, err := s.GetRuntimeAllocation(t.Context(), tenant, session.Environment.ID) + current, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}) if err != nil || current.ObservationError != "resource_missing" { t.Fatal("stale success erased newer failure", current, err) } - if err := w.RecordRuntimeObservation(t.Context(), current, ""); err != nil { + if err := deploymentExecution(t, w).RecordObservation(t.Context(), current, ""); err != nil { t.Fatal(err) } - if err := w.RecordRuntimeObservation(t.Context(), owner, "node_unavailable"); err != nil { + if err := deploymentExecution(t, w).RecordObservation(t.Context(), owner, "node_unavailable"); err != nil { t.Fatal(err) } - recovered, err := s.GetRuntimeAllocation(t.Context(), tenant, session.Environment.ID) + recovered, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}) if err != nil || recovered.ObservationError != "" || recovered.ID != owner.ID { t.Fatal("stale error replaced recovered observation", recovered, err) } - if err := w.RecordRuntimeObservation(t.Context(), current, "secret provider exception"); !errors.Is(err, sessions.ErrInvalidInput) { + if err := deploymentExecution(t, w).RecordObservation(t.Context(), current, "secret provider exception"); !errors.Is(err, deployment.ErrInvalidInput) { t.Fatal("raw diagnostics accepted", err) } } diff --git a/services/core/internal/store/runtime_pending_test.go b/services/core/internal/store/runtime_pending_test.go index f21831945..9a94fefa6 100644 --- a/services/core/internal/store/runtime_pending_test.go +++ b/services/core/internal/store/runtime_pending_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "testing" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "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" @@ -48,15 +49,15 @@ func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) t.Fatal(err) } } - idleOwner, err := s.GetRuntimeAllocation(t.Context(), tenant, idleEnvironment.ID) + idleOwner, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: idleEnvironment.ID}) if err != nil || idleOwner.State != "running" { t.Fatal("idle creation was stranded", idleOwner, err) } - initialOwner, err := s.GetRuntimeAllocation(t.Context(), tenant, initial.Environment.ID) + initialOwner, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: initial.Environment.ID}) if err != nil || initialOwner.State != "running" { t.Fatal("initial creation was stranded", initialOwner, err) } - if _, err := s.GetRuntimeAllocation(t.Context(), deleted.TenantID, deletedEnvironment.ID); err == nil { + if _, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: deleted.TenantID, EnvironmentID: deletedEnvironment.ID}); err == nil { t.Fatal("deleted Session provisioned") } waiting, err := s.GetSession(t.Context(), tenant, initial.ID) @@ -80,8 +81,8 @@ func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) if p.creates != creates { t.Fatal("restart repeated bootstrap", creates, p.creates) } - for _, owner := range []store.RuntimeAllocation{idleOwner, initialOwner} { - got, err := s.GetRuntimeAllocation(t.Context(), tenant, owner.EnvironmentID) + for _, owner := range []deployment.Allocation{idleOwner, initialOwner} { + got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: owner.EnvironmentID}) if err != nil || got.ID != owner.ID || got.DeviceID != owner.DeviceID { t.Fatal("restart replaced allocation identity", got, err) } diff --git a/services/core/internal/store/runtime_placements.go b/services/core/internal/store/runtime_placements.go deleted file mode 100644 index bf35cbf3f..000000000 --- a/services/core/internal/store/runtime_placements.go +++ /dev/null @@ -1,127 +0,0 @@ -package store - -import ( - "context" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/jackc/pgx/v5/pgtype" -) - -// reserveRuntimePlacement shares the deployment lock with restore and node removal. -// A Session creation retry never reaches this function. A node enrolled with -// another Core address receives no new sandbox; restores still reach it. -func reserveRuntimePlacement(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, publicURL string) error { - d, err := q.LockRuntimeDeployment(ctx) - if err != nil { - return err - } - if d.ResetClear.Valid { - return ErrSandboxResetAdmission - } - // A changed installation address cannot admit guests that require a public - // origin. Existing owned resources remain available for cleanup. - if d.ProviderKind != "" { - publicOrigin, err := sandboxProviders.RequiresPublicOrigin(d.ProviderKind) - if err != nil { - return err - } - if publicOrigin && deployment.LoopbackOrigin(publicURL) { - return deployment.ErrPublicURLUnreachable - } - } - - if d.Mode == "direct" { - if d.AdmissionPaused { - return deployment.ErrNodeUnavailable - } - - return nil - } - if d.ProviderKind == "" { - if d.WebManaged { - return deployment.ErrNodeUnavailable - } - return nil - } - if d.AdmissionPaused { - return deployment.ErrNodeUnavailable - } - rows, err := q.ListRuntimeNodes(ctx, pgtype.UUID{}) - if err != nil { - return err - } - var chosen *sqlc.ListRuntimeNodesRow - preparing := false - for i := range rows { - n := &rows[i] - if n.Online && n.TargetState == "preparing" && n.Active < int64(n.MaxActive) && n.Retained < int64(n.MaxRetained) && (!d.WebManaged || n.CoreUrl == publicURL) { - preparing = true - } - if !n.Online || !n.ServingReady || !n.ReadyGeneration.Valid || n.Active >= int64(n.MaxActive) || n.Retained >= int64(n.MaxRetained) || (d.WebManaged && n.CoreUrl != publicURL) { - continue - } - if chosen == nil || n.ReadyGeneration.Int64 > chosen.ReadyGeneration.Int64 || n.ReadyGeneration.Int64 == chosen.ReadyGeneration.Int64 && n.Active < chosen.Active { - chosen = n - } - } - if chosen == nil { - if preparing { - return deployment.ErrNodesPreparing - } - return deployment.ErrNodeUnavailable - } - return q.CreateSessionRuntimePlacement(ctx, sqlc.CreateSessionRuntimePlacementParams{SessionID: session, NodeID: chosen.ID, Generation: chosen.ReadyGeneration.Int64}) -} - -// ResolveRuntimeNode includes deleted Sessions so owned cleanup remains routable. -func (s *Store) ResolveRuntimeNode(ctx context.Context, tenant, environment string) (string, error) { - allocation, err := s.GetRuntimeAllocation(ctx, tenant, environment) - if err != nil { - return "", err - } - if allocation.NodeID == "" { - return "", deployment.ErrNodeUnavailable - } - return allocation.NodeID, nil -} -func reserveRuntimeRestore(ctx context.Context, q *sqlc.Queries, node pgtype.UUID, generation pgtype.Int8) error { - if !node.Valid { - return nil - } - if _, err := q.LockRuntimeDeployment(ctx); err != nil { - return err - } - nodes, err := q.ListRuntimeNodes(ctx, pgtype.UUID{}) - if err != nil { - return err - } - for _, n := range nodes { - if n.ID == node { - ready, err := q.NodeGenerationReady(ctx, sqlc.NodeGenerationReadyParams{NodeID: node, Generation: generation.Int64}) - if err != nil { - return err - } - if !generation.Valid || !n.Online || !ready || n.Active >= int64(n.MaxActive) { - return deployment.ErrNodeUnavailable - } - return nil - } - } - return deployment.ErrNodeUnavailable -} - -// ResolveRuntimeGeneration includes deleted Sessions and never substitutes the target. -func (s *Store) ResolveRuntimeGeneration(ctx context.Context, ref sandbox.Reference) (string, uint64, error) { - a, err := s.GetRuntimeAllocation(ctx, ref.TenantID, ref.EnvironmentID) - if err != nil { - return "", 0, err - } - if a.State == "released" || a.ID != ref.AllocationID || a.NodeID == "" || a.DeploymentGeneration == 0 { - return "", 0, deployment.ErrNodeUnavailable - } - return a.NodeID, a.DeploymentGeneration, nil -} diff --git a/services/core/internal/store/runtime_scan_test.go b/services/core/internal/store/runtime_scan_test.go index 4845c2b03..fe2305218 100644 --- a/services/core/internal/store/runtime_scan_test.go +++ b/services/core/internal/store/runtime_scan_test.go @@ -10,6 +10,7 @@ import ( "github.com/google/uuid" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" ) @@ -94,7 +95,7 @@ func TestManagedRuntimeScanEmptyAfterCleanupAndCanceledCall(t *testing.T) { if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - got, err := s.GetRuntimeAllocation(t.Context(), tenant, env.ID) + got, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}) if err != nil || got.State != "released" || p.kills != 1 { t.Fatal("cleanup delayed at EOF", got, err, p.kills) } diff --git a/services/core/internal/store/runtime_suspension.go b/services/core/internal/store/runtime_suspension.go deleted file mode 100644 index da67bd122..000000000 --- a/services/core/internal/store/runtime_suspension.go +++ /dev/null @@ -1,174 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -// RuntimeActivity separates real work from the connection keepalive. -type RuntimeActivity struct { - LastActivity, ObservedAt time.Time - Busy, WakeRequested, HasCompletedTurn bool -} - -// ReadyToSuspend evaluates the idle policy using one database-clock observation. -func (a RuntimeActivity) ReadyToSuspend(idleTimeout time.Duration) bool { - return idleTimeout > 0 && a.HasCompletedTurn && !a.Busy && !a.WakeRequested && - a.ObservedAt.Sub(a.LastActivity) >= idleTimeout -} - -func runtimeActivity(row sqlc.GetRuntimeActivityRow) RuntimeActivity { - return RuntimeActivity{LastActivity: row.LastActivity.Time, ObservedAt: row.ObservedAt.Time, Busy: row.Busy, WakeRequested: row.ComputeWakeRequested, HasCompletedTurn: row.HasCompletedTurn} -} - -// SetRuntimeCompute commits an operation phase before its external effects. -// Revision and the existing Session lock fence a stale lifecycle observation. -func (s *Store) SetRuntimeCompute(ctx context.Context, owner RuntimeAllocation, phase string, state json.RawMessage, retainedUntil *time.Time, idleTimeout time.Duration) (RuntimeAllocation, error) { - if !runtimeComputeTransition(owner.ComputePhase, phase) || !json.Valid(state) || (owner.ComputePhase == "running" && phase == "quiescing" && idleTimeout <= 0) { - return RuntimeAllocation{}, sessions.ErrInvalidInput - } - if phase != "running" && (retainedUntil == nil || retainedUntil.IsZero()) { - return RuntimeAllocation{}, sessions.ErrInvalidInput - } - var object map[string]json.RawMessage - if json.Unmarshal(state, &object) != nil || object == nil { - return RuntimeAllocation{}, sessions.ErrInvalidInput - } - return s.mutateRuntimeAllocation(ctx, owner, true, func(ctx context.Context, q *sqlc.Queries, row sqlc.RuntimeAllocation) (sqlc.RuntimeAllocation, error) { - if row.ComputeRevision != owner.ComputeRevision || row.ComputePhase != owner.ComputePhase { - return sqlc.RuntimeAllocation{}, sessions.ErrTurnConflict - } - if (phase == "quiescing" && row.ComputePhase == "running") || (phase == "suspending" && row.ComputePhase == "quiescing") { - activity, err := q.GetRuntimeActivity(ctx, row.ID) - if err != nil { - return sqlc.RuntimeAllocation{}, err - } - if activity.Busy || activity.ComputeWakeRequested { - return sqlc.RuntimeAllocation{}, sessions.ErrTurnConflict - } - if phase == "quiescing" && (!runtimeActivity(activity).ReadyToSuspend(idleTimeout) || row.ComputeActivityAt.Time.After(owner.ComputeActivityAt)) { - return sqlc.RuntimeAllocation{}, sessions.ErrTurnConflict - } - } - if row.ComputePhase == "suspended" && phase == "restoring" { - if err := reserveRuntimeRestore(ctx, q, row.NodeID, row.DeploymentGeneration); err != nil { - return sqlc.RuntimeAllocation{}, err - } - } - until := pgtype.Timestamptz{} - if retainedUntil != nil { - until = pgtype.Timestamptz{Time: *retainedUntil, Valid: true} - } - return q.SetRuntimeCompute(ctx, sqlc.SetRuntimeComputeParams{ - ID: row.ID, Revision: row.ComputeRevision, Phase: phase, State: state, RetainedUntil: until, - }) - }) -} - -func runtimeComputeTransition(from, to string) bool { - if from == to { - return from != "disabled" - } - switch from { - case "disabled": - return to == "running" - case "running": - return to == "quiescing" - case "quiescing": - return to == "suspending" || to == "waking" || to == "running" - case "suspending": - return to == "suspended" || to == "waking" - case "suspended": - return to == "restoring" - case "restoring": - return to == "waking" - case "waking": - return to == "running" - } - return false -} - -func (s *Store) RuntimeActivity(ctx context.Context, owner RuntimeAllocation) (RuntimeActivity, error) { - id, err := parseID(owner.ID) - if err != nil { - return RuntimeActivity{}, err - } - if err := s.checkExecutionOwnership(ctx); err != nil { - return RuntimeActivity{}, err - } - row, err := s.queries.GetRuntimeActivity(ctx, id) - if errors.Is(err, pgx.ErrNoRows) { - return RuntimeActivity{}, sessions.ErrNotFound - } - if err != nil { - return RuntimeActivity{}, err - } - return runtimeActivity(row), nil -} - -// TouchRuntimeActivity is used only by operations requiring live compute. -// Public history and published-artifact reads do not call it. -func (s *Store) TouchRuntimeActivity(ctx context.Context, tenant, environment string) error { - lookup, err := sessionpg.DeviceLookup(tenant, environment) - if err != nil { - return err - } - owned, err := sessionpg.LoadEnvironment(ctx, s.queries, tenant, environment) - if err != nil { - return err - } - return s.withPublicSession(ctx, tenant, owned.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - current, err := q.GetEnvironment(ctx, sqlc.GetEnvironmentParams{TenantID: lookup.TenantID, ID: lookup.ID}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - if current.Environment.SessionID != session { - return sessions.ErrNotFound - } - return q.TouchRuntimeActivity(ctx, sqlc.TouchRuntimeActivityParams{TenantID: lookup.TenantID, EnvironmentID: lookup.ID}) - }) -} - -func (s *Store) ClearRuntimeWake(ctx context.Context, owner RuntimeAllocation, observedActivity time.Time) error { - id, err := parseID(owner.ID) - if err != nil { - return err - } - 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}}) -} - -func (s *Store) CountRuntimeComputeReservations(ctx context.Context, provider string) (int64, error) { - id, err := parseID(provider) - if err != nil { - return 0, err - } - if err := s.checkExecutionOwnership(ctx); err != nil { - return 0, err - } - return s.queries.CountRuntimeComputeReservations(ctx, id) -} - -func (s *Store) CountRuntimeRetainedAllocations(ctx context.Context, provider string) (int64, error) { - id, err := parseID(provider) - if err != nil { - return 0, err - } - 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_concurrency_test.go b/services/core/internal/store/runtime_suspension_concurrency_test.go index 09f8d9f9a..5d2a97fd7 100644 --- a/services/core/internal/store/runtime_suspension_concurrency_test.go +++ b/services/core/internal/store/runtime_suspension_concurrency_test.go @@ -8,6 +8,7 @@ import ( "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" @@ -123,7 +124,7 @@ func TestRuntimeSuspensionCaptureRechecksNewPendingWork(t *testing.T) { ctx, tx, blocker := runtimeSuspensionLockedSession(t, pool, owner.SessionID) done := make(chan error, 1) go func() { - _, err := w.SetRuntimeCompute(ctx, owner, "suspending", json.RawMessage(`{}`), &until, 0) + _, err := deploymentExecution(t, w).SetCompute(ctx, owner, "suspending", json.RawMessage(`{}`), &until, 0) done <- err }() runtimeSuspensionWaitBlocked(t, ctx, pool, blocker, done) @@ -144,10 +145,10 @@ func TestRuntimeSuspensionCaptureRechecksNewPendingWork(t *testing.T) { if err := tx.Commit(ctx); err != nil { t.Fatal(err) } - if err := <-done; !errors.Is(err, sessions.ErrTurnConflict) { + if err := <-done; !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("capture ignored work admitted after quiesce", err) } - got, err := s.GetRuntimeAllocation(ctx, owner.TenantID, owner.EnvironmentID) + got, err := deploymentStore(s).EnvironmentAllocation(ctx, deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil || got.ComputePhase != "quiescing" || got.ComputeRevision != owner.ComputeRevision { t.Fatal("failed capture CAS changed its owner", got, err) } @@ -161,7 +162,7 @@ func TestRuntimeSuspensionWakeUsesSessionLock(t *testing.T) { s, _, pool, owner := runtimeSuspensionFixture(t) ctx, tx, blocker := runtimeSuspensionLockedSession(t, pool, owner.SessionID) done := make(chan error, 1) - go func() { done <- s.TouchRuntimeActivity(ctx, owner.TenantID, owner.EnvironmentID) }() + go func() { done <- deploymentService(t, s).TouchActivity(ctx, owner.TenantID, owner.EnvironmentID) }() runtimeSuspensionWaitBlocked(t, ctx, pool, blocker, done) if deleted { if _, err := tx.Exec(ctx, `UPDATE sessions SET deleted_at=clock_timestamp() WHERE id=$1`, owner.SessionID); err != nil { @@ -177,7 +178,7 @@ func TestRuntimeSuspensionWakeUsesSessionLock(t *testing.T) { if deleted && !errors.Is(err, sessions.ErrNotFound) || !deleted && err != nil { t.Fatal("wake did not observe locked state", err) } - got, err := s.GetRuntimeAllocation(ctx, owner.TenantID, owner.EnvironmentID) + got, err := deploymentStore(s).EnvironmentAllocation(ctx, deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil || got.ComputeWakeRequested == deleted { t.Fatal("wake was lost or crossed deletion", got, err) } @@ -192,7 +193,7 @@ func TestRuntimeSuspensionQuiesceCannotOvertakeClaim(t *testing.T) { ctx, tx, blocker := runtimeSuspensionLockedSession(t, pool, owner.SessionID) done := make(chan error, 1) go func() { - _, err := w.SetRuntimeCompute(ctx, owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond) + _, err := deploymentExecution(t, w).SetCompute(ctx, owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond) done <- err }() runtimeSuspensionWaitBlocked(t, ctx, pool, blocker, done) @@ -210,10 +211,10 @@ func TestRuntimeSuspensionQuiesceCannotOvertakeClaim(t *testing.T) { if err := tx.Commit(ctx); err != nil { t.Fatal(err) } - if err := <-done; !errors.Is(err, sessions.ErrTurnConflict) { + if err := <-done; !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("quiesce overtook an admitted Turn", err) } - got, err := s.GetRuntimeAllocation(ctx, owner.TenantID, owner.EnvironmentID) + got, err := deploymentStore(s).EnvironmentAllocation(ctx, deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil || got.ComputePhase != "running" || got.ComputeRevision != owner.ComputeRevision { t.Fatal("active work was quiesced", got, err) } diff --git a/services/core/internal/store/runtime_suspension_test.go b/services/core/internal/store/runtime_suspension_test.go index df515534e..d768c9597 100644 --- a/services/core/internal/store/runtime_suspension_test.go +++ b/services/core/internal/store/runtime_suspension_test.go @@ -8,27 +8,28 @@ import ( "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" ) -func runtimeSuspensionFixture(t *testing.T) (*Store, *Store, *pgxpool.Pool, RuntimeAllocation) { +func runtimeSuspensionFixture(t *testing.T) (*Store, *Store, *pgxpool.Pool, deployment.Allocation) { t.Helper() s, pool := testStore(t) w := executionWriter(t, s) tenant := uuid.NewString() _, environment := localEnvironment(t, s, tenant) - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - owner, err = w.ObserveRuntimeRunning(t.Context(), owner) + owner, err = deploymentExecution(t, w).ObserveRunning(t.Context(), owner) if err != nil { t.Fatal(err) } - owner, err = w.SetRuntimeCompute(t.Context(), owner, "running", json.RawMessage(`{"instance":"original"}`), nil, 0) + owner, err = deploymentExecution(t, w).SetCompute(t.Context(), owner, "running", json.RawMessage(`{"instance":"original"}`), nil, 0) if err != nil { t.Fatal(err) } @@ -42,20 +43,20 @@ func runtimeSuspensionSQL(t *testing.T, pool *pgxpool.Pool, query string, args . } } -func runtimeSuspensionCompleted(t *testing.T, pool *pgxpool.Pool, owner RuntimeAllocation) string { +func runtimeSuspensionCompleted(t *testing.T, pool *pgxpool.Pool, owner deployment.Allocation) string { t.Helper() id := uuid.NewString() runtimeSuspensionSQL(t, pool, `INSERT INTO turns(id,session_id,status,completed_at) VALUES($1,$2,'completed',clock_timestamp()-interval '10 minutes')`, id, owner.SessionID) return id } -func runtimeSuspensionStep(t *testing.T, w *Store, owner RuntimeAllocation, phase string, until *time.Time) RuntimeAllocation { +func runtimeSuspensionStep(t *testing.T, w *Store, owner deployment.Allocation, phase string, until *time.Time) deployment.Allocation { t.Helper() idleTimeout := time.Duration(0) if owner.ComputePhase == "running" && phase == "quiescing" { idleTimeout = time.Nanosecond } - next, err := w.SetRuntimeCompute(t.Context(), owner, phase, json.RawMessage(`{"instance":"original","snapshot":"qualified"}`), until, idleTimeout) + next, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, phase, json.RawMessage(`{"instance":"original","snapshot":"qualified"}`), until, idleTimeout) if err != nil { t.Fatalf("%s -> %s: %v", owner.ComputePhase, phase, err) } @@ -84,7 +85,7 @@ func TestRuntimeSuspensionRequiresCompletedIdleAndNoPendingWork(t *testing.T) { case "file_write": runtimeSuspensionSQL(t, pool, `INSERT INTO environment_file_writes(id,environment_id,device_id,request_sha256) VALUES($1,$2,$3,$4)`, uuid.NewString(), owner.EnvironmentID, owner.DeviceID, strings.Repeat("a", 64)) } - activity, err := w.RuntimeActivity(t.Context(), owner) + activity, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil { t.Fatal(err) } @@ -93,12 +94,12 @@ func TestRuntimeSuspensionRequiresCompletedIdleAndNoPendingWork(t *testing.T) { t.Fatalf("activity lost pending work: %+v", activity) } until := time.Now().Add(time.Hour) - _, err = w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond) + _, err = deploymentExecution(t, w).SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond) if kind == "idle" { if err != nil { t.Fatal(err) } - } else if !errors.Is(err, sessions.ErrTurnConflict) { + } else if !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatalf("unsafe quiesce admitted: %v", err) } }) @@ -115,7 +116,7 @@ func TestRuntimeSuspensionCASAndActivityFence(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - _, err := w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{"operation":"same-observation"}`), &until, time.Nanosecond) + _, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{"operation":"same-observation"}`), &until, time.Nanosecond) results <- err }() } @@ -125,7 +126,7 @@ func TestRuntimeSuspensionCASAndActivityFence(t *testing.T) { for err := range results { if err == nil { winners++ - } else if errors.Is(err, sessions.ErrTurnConflict) { + } else if errors.Is(err, deployment.ErrAllocationConflict) { conflicts++ } else { t.Fatal(err) @@ -134,38 +135,38 @@ func TestRuntimeSuspensionCASAndActivityFence(t *testing.T) { if winners != 1 || conflicts != 1 { t.Fatalf("CAS admitted competing owners: wins=%d conflicts=%d", winners, conflicts) } - current, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + current, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil { t.Fatal(err) } if current.ComputeRevision != owner.ComputeRevision+1 || current.ComputePhase != "quiescing" { t.Fatal("operation intent not durable", current) } - if _, err := w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{"stale":true}`), &until, time.Nanosecond); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{"stale":true}`), &until, time.Nanosecond); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("stale phase overwrite", err) } current = runtimeSuspensionStep(t, w, current, "running", nil) observed := current - if err := s.TouchRuntimeActivity(t.Context(), owner.TenantID, owner.EnvironmentID); err != nil { + if err := deploymentService(t, s).TouchActivity(t.Context(), owner.TenantID, owner.EnvironmentID); err != nil { t.Fatal(err) } // Clearing an observed wake cannot let an older observation authorize sleep. - latest, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + latest, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil { t.Fatal(err) } - if err := w.ClearRuntimeWake(t.Context(), latest, latest.ComputeActivityAt); err != nil { + if err := deploymentExecution(t, w).ClearWake(t.Context(), latest, latest.ComputeActivityAt); err != nil { t.Fatal(err) } - if _, err := w.SetRuntimeCompute(t.Context(), observed, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), observed, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("newer activity was swallowed", err) } for _, invalid := range []json.RawMessage{json.RawMessage(`[]`), json.RawMessage(`null`), json.RawMessage(`false`), json.RawMessage(`{`)} { - if _, err := w.SetRuntimeCompute(t.Context(), latest, "running", invalid, nil, 0); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), latest, "running", invalid, nil, 0); !errors.Is(err, deployment.ErrInvalidInput) { t.Fatal("non-object compute state accepted", string(invalid), err) } } - if _, err := w.SetRuntimeCompute(t.Context(), latest, "suspended", json.RawMessage(`{}`), &until, 0); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), latest, "suspended", json.RawMessage(`{}`), &until, 0); !errors.Is(err, deployment.ErrInvalidInput) { t.Fatal("running skipped snapshot protocol", err) } } @@ -173,52 +174,52 @@ func TestRuntimeSuspensionCASAndActivityFence(t *testing.T) { func TestRuntimeSuspensionWakeDoesNotLoseNewerWork(t *testing.T) { s, w, pool, owner := runtimeSuspensionFixture(t) runtimeSuspensionCompleted(t, pool, owner) - before, err := w.RuntimeActivity(t.Context(), owner) + before, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil { t.Fatal(err) } - if _, err := w.KeepRuntimeAllocation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).KeepAllocation(t.Context(), owner); err != nil { t.Fatal(err) } if _, err := s.GetSession(t.Context(), owner.TenantID, owner.SessionID); err != nil { t.Fatal(err) } - quiet, err := w.RuntimeActivity(t.Context(), owner) + quiet, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || quiet.WakeRequested || !quiet.LastActivity.Equal(before.LastActivity) { t.Fatal("heartbeat or history read touched compute activity", quiet, err) } - if err := s.TouchRuntimeActivity(t.Context(), uuid.NewString(), owner.EnvironmentID); !errors.Is(err, sessions.ErrNotFound) { + if err := deploymentService(t, s).TouchActivity(t.Context(), uuid.NewString(), owner.EnvironmentID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - foreign, err := w.RuntimeActivity(t.Context(), owner) + foreign, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || foreign.WakeRequested { t.Fatal("foreign tenant woke environment", foreign, err) } - if err := s.TouchRuntimeActivity(t.Context(), owner.TenantID, owner.EnvironmentID); err != nil { + if err := deploymentService(t, s).TouchActivity(t.Context(), owner.TenantID, owner.EnvironmentID); err != nil { t.Fatal(err) } - first, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + first, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil { t.Fatal(err) } - if err := s.TouchRuntimeActivity(t.Context(), owner.TenantID, owner.EnvironmentID); err != nil { + if err := deploymentService(t, s).TouchActivity(t.Context(), owner.TenantID, owner.EnvironmentID); err != nil { t.Fatal(err) } - second, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + second, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil || !second.ComputeActivityAt.After(first.ComputeActivityAt) { t.Fatal("Touch did not record newer activity", err) } - if err := w.ClearRuntimeWake(t.Context(), first, first.ComputeActivityAt); err != nil { + if err := deploymentExecution(t, w).ClearWake(t.Context(), first, first.ComputeActivityAt); err != nil { t.Fatal(err) } - activity, err := w.RuntimeActivity(t.Context(), owner) + activity, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || !activity.WakeRequested { t.Fatal("old wake receipt swallowed newer work", activity, err) } - if err := w.ClearRuntimeWake(t.Context(), second, second.ComputeActivityAt); err != nil { + if err := deploymentExecution(t, w).ClearWake(t.Context(), second, second.ComputeActivityAt); err != nil { t.Fatal(err) } - activity, err = w.RuntimeActivity(t.Context(), owner) + activity, err = deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || activity.WakeRequested { t.Fatal("current wake receipt did not settle", activity, err) } @@ -232,30 +233,30 @@ func TestRuntimeSuspensionRetentionAndDeletedSession(t *testing.T) { owner = runtimeSuspensionStep(t, w, owner, phase, &until) } runtimeSuspensionSQL(t, pool, `UPDATE runtime_allocations SET kept_at=clock_timestamp()-interval '2 hours' WHERE id=$1`, owner.ID) - retained, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + retained, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil || retained.Expired { t.Fatal("suspended snapshot expired by disconnected heartbeat", retained, err) } runtimeSuspensionSQL(t, pool, `UPDATE runtime_allocations SET compute_retained_until=clock_timestamp()-interval '1 second' WHERE id=$1`, owner.ID) - expired, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + expired, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil || !expired.Expired { t.Fatal("snapshot retention expiry not observed", expired, err) } // Use the earlier unexpired observation to exercise expiry at the database CAS. - if _, err := w.SetRuntimeCompute(t.Context(), retained, "restoring", json.RawMessage(`{}`), &until, 0); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), retained, "restoring", json.RawMessage(`{}`), &until, 0); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("expired snapshot restored from stale observation", err) } if err := s.DeleteSession(t.Context(), owner.TenantID, owner.SessionID); err != nil { t.Fatal(err) } - if err := s.TouchRuntimeActivity(t.Context(), owner.TenantID, owner.EnvironmentID); !errors.Is(err, sessions.ErrNotFound) { + if err := deploymentService(t, s).TouchActivity(t.Context(), owner.TenantID, owner.EnvironmentID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - deleted, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + deleted, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil || !deleted.SessionDeleted || deleted.ComputeWakeRequested { t.Fatal("deleted session was woken", deleted, err) } - if _, err := w.SetRuntimeCompute(t.Context(), deleted, "restoring", json.RawMessage(`{}`), &until, 0); !errors.Is(err, sessions.ErrNotFound) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), deleted, "restoring", json.RawMessage(`{}`), &until, 0); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted session restored", err) } } @@ -275,7 +276,7 @@ func TestRuntimeSuspensionCountsUncertainCapacityUntilReleased(t *testing.T) { for _, item := range cases { tenant := uuid.NewString() _, env := localEnvironment(t, s, tenant) - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, env.ID, provider, runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: env.ID}, provider, runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } @@ -286,16 +287,16 @@ func TestRuntimeSuspensionCountsUncertainCapacityUntilReleased(t *testing.T) { if item.state != "released" { wantRetained++ } - retained, err := w.CountRuntimeRetainedAllocations(t.Context(), provider) + retained, err := deploymentStore(w).CountRetainedAllocations(t.Context(), provider) if err != nil || retained != wantRetained { t.Fatalf("retained capacity state=%s phase=%s got=%d want=%d err=%v", item.state, item.phase, retained, wantRetained, err) } - got, err := w.CountRuntimeComputeReservations(t.Context(), provider) + got, err := deploymentStore(w).CountComputeReservations(t.Context(), provider) if err != nil || got != want { t.Fatalf("capacity state=%s phase=%s got=%d want=%d err=%v", item.state, item.phase, got, want, err) } } - got, err := w.CountRuntimeComputeReservations(t.Context(), uuid.NewString()) + got, err := deploymentStore(w).CountComputeReservations(t.Context(), uuid.NewString()) if err != nil || got != 0 { t.Fatal("capacity crossed installation boundary", got, err) } @@ -309,7 +310,7 @@ func TestRuntimeSuspensionIdleStartsAfterLastCompletion(t *testing.T) { if err := pool.QueryRow(t.Context(), `SELECT completed_at FROM turns WHERE id=$1`, turn).Scan(&completed); err != nil { t.Fatal(err) } - activity, err := w.RuntimeActivity(t.Context(), owner) + activity, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || !activity.LastActivity.Equal(completed) { t.Fatal("long Turn completion did not restart idle interval", activity, completed, err) } @@ -322,17 +323,17 @@ func TestRuntimeSuspensionWakeRemainsUntilRunning(t *testing.T) { for _, phase := range []string{"quiescing", "suspending", "suspended"} { owner = runtimeSuspensionStep(t, w, owner, phase, &until) } - if err := s.TouchRuntimeActivity(t.Context(), owner.TenantID, owner.EnvironmentID); err != nil { + if err := deploymentService(t, s).TouchActivity(t.Context(), owner.TenantID, owner.EnvironmentID); err != nil { t.Fatal(err) } - observed, err := s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + observed, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil { t.Fatal(err) } - if err := w.ClearRuntimeWake(t.Context(), owner, observed.ComputeActivityAt); err != nil { + if err := deploymentExecution(t, w).ClearWake(t.Context(), owner, observed.ComputeActivityAt); err != nil { t.Fatal(err) } - activity, err := w.RuntimeActivity(t.Context(), owner) + activity, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || !activity.WakeRequested { t.Fatal("wake cleared before compute was running", activity, err) } @@ -343,10 +344,10 @@ func TestRuntimeSuspensionWakeRemainsUntilRunning(t *testing.T) { if owner.ComputeRetainedUntil != nil { t.Fatal("running retained snapshot expiration") } - if err := w.ClearRuntimeWake(t.Context(), owner, observed.ComputeActivityAt); err != nil { + if err := deploymentExecution(t, w).ClearWake(t.Context(), owner, observed.ComputeActivityAt); err != nil { t.Fatal(err) } - activity, err = w.RuntimeActivity(t.Context(), owner) + activity, err = deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || activity.WakeRequested { t.Fatal("running compute could not settle wake", activity, err) } @@ -357,24 +358,15 @@ func TestRuntimeSuspensionExpiredRunningAndLostWriterAreFenced(t *testing.T) { runtimeSuspensionCompleted(t, pool, owner) until := time.Now().Add(time.Hour) runtimeSuspensionSQL(t, pool, `UPDATE runtime_allocations SET kept_at=clock_timestamp()-interval '2 hours' WHERE id=$1`, owner.ID) - if _, err := w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("expired running allocation entered checkpoint", err) } if err := w.lease.Close(t.Context()); err != nil { t.Fatal(err) } - if _, err := w.RuntimeActivity(t.Context(), owner); err == nil { - t.Fatal("lost writer read execution activity") - } - if err := w.ClearRuntimeWake(t.Context(), owner, owner.ComputeActivityAt); err == nil { + if err := deploymentExecution(t, w).ClearWake(t.Context(), owner, owner.ComputeActivityAt); err == nil { t.Fatal("lost writer changed wake receipt") } - if _, err := w.CountRuntimeComputeReservations(t.Context(), owner.ProviderKey); err == nil { - t.Fatal("lost writer admitted running capacity") - } - if _, err := w.CountRuntimeRetainedAllocations(t.Context(), owner.ProviderKey); err == nil { - t.Fatal("lost writer admitted retained capacity") - } } func TestRuntimeSuspensionRechecksCompletionAgainstIdleTimeout(t *testing.T) { @@ -384,12 +376,12 @@ func TestRuntimeSuspensionRechecksCompletionAgainstIdleTimeout(t *testing.T) { turn := runtimeSuspensionCompleted(t, pool, owner) runtimeSuspensionSQL(t, pool, `UPDATE runtime_allocations SET compute_activity_at=clock_timestamp()-interval '20 minutes' WHERE id=$1`, owner.ID) var err error - owner, err = s.GetRuntimeAllocation(t.Context(), owner.TenantID, owner.EnvironmentID) + owner, err = deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: owner.TenantID, EnvironmentID: owner.EnvironmentID}) if err != nil { t.Fatal(err) } idleTimeout := time.Minute - observed, err := w.RuntimeActivity(t.Context(), owner) + observed, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || !observed.ReadyToSuspend(idleTimeout) { t.Fatal("fixture is not initially idle", observed, err) } @@ -407,17 +399,17 @@ func TestRuntimeSuspensionRechecksCompletionAgainstIdleTimeout(t *testing.T) { runtimeSuspensionSQL(t, pool, `INSERT INTO environment_file_writes(id,environment_id,device_id,request_sha256,state,created_at,settled_at) VALUES($1,$2,$3,$4,$5,clock_timestamp()-interval '10 minutes',clock_timestamp())`, uuid.NewString(), owner.EnvironmentID, owner.DeviceID, strings.Repeat("a", 64), strings.TrimPrefix(kind, "file_")) } until := time.Now().Add(time.Hour) - if _, err := w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, idleTimeout); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, idleTimeout); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("completion after idle observation did not fence quiesce", err) } - activity, err := w.RuntimeActivity(t.Context(), owner) + activity, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || activity.Busy || activity.WakeRequested || activity.ReadyToSuspend(idleTimeout) { t.Fatal("last completion did not restart idle interval", activity, err) } - if _, err := w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, 0); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, 0); !errors.Is(err, deployment.ErrInvalidInput) { t.Fatal("missing idle timeout accepted", err) } - if _, err := w.SetRuntimeCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond); err != nil { + if _, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond); err != nil { t.Fatal("elapsed idle timeout rejected", err) } }) @@ -437,7 +429,7 @@ func TestRuntimeComputePhaseChangedAt(t *testing.T) { return value } entered := changedAt() - owner, err := w.SetRuntimeCompute(t.Context(), owner, "running", json.RawMessage(`{"instance":"updated"}`), nil, 0) + owner, err := deploymentExecution(t, w).SetCompute(t.Context(), owner, "running", json.RawMessage(`{"instance":"updated"}`), nil, 0) if err != nil { t.Fatal(err) } @@ -459,13 +451,13 @@ func TestRuntimeComputePhaseChangedAtInNodeAllocations(t *testing.T) { if err != nil { t.Fatal(err) } - allocation, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, d.InstallationID, runtimedevice.HashCredential("runtime")) + allocation, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, d.InstallationID, runtimedevice.HashCredential("runtime")) if err != nil { t.Fatal(err) } - listed := func() RuntimeNodeAllocation { + listed := func() deployment.NodeAllocation { t.Helper() - items, err := s.ListNodeRuntimeAllocations(t.Context(), d.LocalNodeID) + items, err := deploymentStore(s).NodeAllocations(t.Context(), d.LocalNodeID) if err != nil || len(items) != 1 || items[0].ID != allocation.ID { t.Fatal(items, err) } 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 1b61291f6..d6744f12d 100644 --- a/services/core/internal/store/runtime_wake_hint_integration_test.go +++ b/services/core/internal/store/runtime_wake_hint_integration_test.go @@ -9,10 +9,10 @@ import ( "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) type wakeHintScanProvider struct { @@ -42,7 +42,7 @@ type wakeHintIntegrationTarget struct { tenant string session sessions.Session environment sessions.Environment - owner store.RuntimeAllocation + owner deployment.Allocation } type wakeHintIntegration struct { @@ -166,7 +166,7 @@ func TestRuntimeWakeHintCommittedSubmitResumesBeforeNormalTick(t *testing.T) { } f.release() awaitDaemonRemoteCondition(t, t.Context(), 2*time.Second, "hint restored retained compute", func() bool { - owner, err := f.fixture.store.GetRuntimeAllocation(t.Context(), f.target.tenant, f.target.environment.ID) + owner, err := fixtureReader(f.fixture.db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: f.target.tenant, EnvironmentID: f.target.environment.ID}) return err == nil && owner.ComputePhase == "running" }) if elapsed := time.Since(f.started); elapsed >= 3*time.Second { diff --git a/services/core/internal/store/sandbox_credential_allocations.go b/services/core/internal/store/sandbox_credential_allocations.go deleted file mode 100644 index 879dd6929..000000000 --- a/services/core/internal/store/sandbox_credential_allocations.go +++ /dev/null @@ -1,28 +0,0 @@ -package store - -import ( - "context" - - "github.com/jackc/pgx/v5/pgtype" -) - -// SandboxCredentialAllocationPage reads owned receipts without granting execution authority. -func (s *Store) SandboxCredentialAllocationPage(ctx context.Context, after string) ([]RuntimeAllocation, error) { - id := pgtype.UUID{Valid: true} - if after != "" { - var err error - id, err = parseID(after) - if err != nil { - return nil, err - } - } - rows, err := s.queries.ListRuntimeAllocations(ctx, id) - if err != nil { - return nil, err - } - out := make([]RuntimeAllocation, 0, len(rows)) - for _, r := range rows { - out = append(out, runtimeAllocationFromRow(r.RuntimeAllocation, r.SessionID, r.TenantID, r.DeletedAt, r.Expired)) - } - return out, nil -} diff --git a/services/core/internal/store/sandbox_deployment_resources_test.go b/services/core/internal/store/sandbox_deployment_resources_test.go index 6368615db..cd2595cb0 100644 --- a/services/core/internal/store/sandbox_deployment_resources_test.go +++ b/services/core/internal/store/sandbox_deployment_resources_test.go @@ -16,7 +16,7 @@ func TestSandboxDeploymentMutationViewsIncludeActualResources(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) changes := deploymentExecution(t, w) installation := uuid.NewString() @@ -32,7 +32,7 @@ func TestSandboxDeploymentMutationViewsIncludeActualResources(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, installation, runtimedevice.HashCredential(uuid.NewString())); err != nil { + if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, installation, runtimedevice.HashCredential(uuid.NewString())); err != nil { t.Fatal(err) } if _, err := s.CreateSession(t.Context(), tenant, managerSessionInput(uuid.NewString())); err != nil { diff --git a/services/core/internal/store/sandbox_deployment_setup_test.go b/services/core/internal/store/sandbox_deployment_setup_test.go index 5e1c39435..6415add8e 100644 --- a/services/core/internal/store/sandbox_deployment_setup_test.go +++ b/services/core/internal/store/sandbox_deployment_setup_test.go @@ -8,6 +8,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/google/uuid" @@ -15,7 +16,7 @@ import ( func TestSandboxDeploymentSetupPersistsWithoutExecution(t *testing.T) { s, pool := newManagedTestStore(t) - s.SetPublicURL("https://core.example") + s.SetPlacement(placementRules(t, "https://core.example")) w := executionWriter(t, s) id := uuid.NewString() if err := deploymentExecution(t, w).Claim(t.Context(), id); err != nil { @@ -25,7 +26,7 @@ func TestSandboxDeploymentSetupPersistsWithoutExecution(t *testing.T) { if err != nil || before.InstallationID != id || before.Provider != "" || before.CoreURL != "https://core.example" { t.Fatal(before, err) } - if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, deployment.ErrNodeUnavailable) { + if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, placement.ErrNodeUnavailable) { t.Fatal("uninitialized hosted admission", err) } input := sandbox.Selection{DeploymentSpec: SandboxDeploymentTestSpec("microsandbox"), Provider: "microsandbox"} diff --git a/services/core/internal/store/sandbox_deployment_switch_test.go b/services/core/internal/store/sandbox_deployment_switch_test.go index a101ea7f5..3447ee77d 100644 --- a/services/core/internal/store/sandbox_deployment_switch_test.go +++ b/services/core/internal/store/sandbox_deployment_switch_test.go @@ -10,6 +10,7 @@ import ( "sync" "testing" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/e2b" @@ -36,7 +37,7 @@ func TestSandboxResetClearsCustomE2BEndpoint(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) changes := deploymentExecution(t, w) installation := uuid.NewString() @@ -77,7 +78,7 @@ func TestSandboxDirectDeploymentOwnershipAndCleanSwitch(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) changes := deploymentExecution(t, w) id := uuid.NewString() @@ -110,11 +111,11 @@ func TestSandboxDirectDeploymentOwnershipAndCleanSwitch(t *testing.T) { t.Fatal(err) } environment := session.Environment.ID - nodes, err := w.ListRuntimeLifecycleNodes(t.Context()) + nodes, err := deploymentStore(w).LifecycleNodes(t.Context()) if err != nil || len(nodes) != 1 || nodes[0] != "" { t.Fatal("cloud lifecycle requires node", nodes, err) } - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment, id, runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}, id, runtimedevice.HashCredential(uuid.NewString())) if err != nil || owner.NodeID != "" { t.Fatal(owner, err) } @@ -125,7 +126,7 @@ func TestSandboxDirectDeploymentOwnershipAndCleanSwitch(t *testing.T) { if _, err := pool.Exec(t.Context(), "UPDATE runtime_allocations SET kept_at=clock_timestamp()-interval '2 hours' WHERE id=$1", owner.ID); err != nil { t.Fatal(err) } - observed, err := w.GetRuntimeAllocation(t.Context(), tenant, environment) + observed, err := deploymentStore(w).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment}) if err != nil || observed.Expired { t.Fatal("cloud inherited legacy node-less expiry", err) } @@ -136,16 +137,16 @@ func TestSandboxDirectDeploymentOwnershipAndCleanSwitch(t *testing.T) { if _, err := changes.Update(SandboxResetTestContext(t.Context()), id, update); !errors.Is(err, deployment.ErrResetInProgress) { t.Fatal("reset allowed switch", err) } - if _, err := w.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), owner); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("unsettled create released", err) } - if _, err := w.SettleRuntimeCreation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).SettleCreation(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } changed, err := resetAndSelect(t, w, id, update.ExpectedGeneration, update) @@ -163,7 +164,7 @@ func TestSandboxDirectDeploymentOwnershipAndCleanSwitch(t *testing.T) { func TestSandboxSwitchRetiresNodesAndEnrollment(t *testing.T) { _, pool := newManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{5}, 32)) - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) changes := deploymentExecution(t, w) nodes := deploymentService(t, s) @@ -238,7 +239,7 @@ func TestSandboxSwitchRetiresNodesAndEnrollment(t *testing.T) { func TestSandboxResetSerializesFreshDirectSessions(t *testing.T) { _, pool := newManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{6}, 32)) - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) changes := deploymentExecution(t, w) id := uuid.NewString() @@ -255,7 +256,7 @@ func TestSandboxResetSerializesFreshDirectSessions(t *testing.T) { go func() { defer wg.Done() _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())) - if err != nil && !errors.Is(err, ErrSandboxResetAdmission) && !errors.Is(err, deployment.ErrNodeUnavailable) { + if err != nil && !errors.Is(err, placement.ErrResetAdmission) && !errors.Is(err, placement.ErrNodeUnavailable) { t.Error(err) } }() @@ -265,7 +266,7 @@ func TestSandboxResetSerializesFreshDirectSessions(t *testing.T) { } wg.Wait() for range 3 { - if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, ErrSandboxResetAdmission) { + if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, placement.ErrResetAdmission) { t.Fatal("fresh creation bypassed reset", err) } } @@ -282,7 +283,7 @@ func TestSandboxResetSerializesFreshDirectSessions(t *testing.T) { func TestSandboxSwitchPreservesReleasedAllocationAndItemHistory(t *testing.T) { _, pool := newManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{8}, 32)) - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) changes := deploymentExecution(t, w) installation := uuid.NewString() @@ -298,11 +299,11 @@ func TestSandboxSwitchPreservesReleasedAllocationAndItemHistory(t *testing.T) { if err != nil { t.Fatal(err) } - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, installation, runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, installation, runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - if _, err := w.ReleaseAbsentRuntimeCreation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).ReleaseAbsentCreation(t.Context(), owner); err != nil { t.Fatal(err) } // A separate completed Session supplies public Items without calling a model. @@ -367,7 +368,7 @@ func TestSandboxSwitchPreservesReleasedAllocationAndItemHistory(t *testing.T) { func TestUnspecifiedNodeDeploymentRejectedWithoutMutation(t *testing.T) { _, pool := newManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{7}, 32)) - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) changes := deploymentExecution(t, w) nodes := deploymentService(t, s) @@ -404,7 +405,7 @@ func TestUnspecifiedNodeDeploymentRejectedWithoutMutation(t *testing.T) { if _, err := nodes.NodeConfiguration(t.Context(), node.NodeID, node.Credential, 0); !errors.Is(err, deployment.ErrSpecificationMismatch) { t.Fatal("node configuration served without a specification", err) } - if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, sessions.ErrEnvironmentUnavailable) { + if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, placement.ErrAdmissionClosed) { t.Fatal("unspecified deployment admitted a fresh sandbox", err) } if _, err := nodes.CreateEnrollment(t.Context(), deployment.Capacity{MaxActive: 2, MaxRetained: 4}); !errors.Is(err, deployment.ErrConflict) { 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 c9c8a5021..64d6118ff 100644 --- a/services/core/internal/store/sandbox_deployment_switch_worker_test.go +++ b/services/core/internal/store/sandbox_deployment_switch_worker_test.go @@ -25,6 +25,7 @@ func TestSandboxWorkerSwitchesAndRecoversFailedActivation(t *testing.T) { _, pool := store.NewManagedTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{7}, 32)) s, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + s.SetPlacement(fixtureRules(t, db)) deployments := fixtureDeployment(t, db) id := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} @@ -157,7 +158,7 @@ func TestSandboxWorkerSwitchesAndRecoversFailedActivation(t *testing.T) { if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { t.Fatal(err) } - reconcileManagedState(t, w, s, tenant, environment.ID, "released") + reconcileManagedState(t, w, db, tenant, environment.ID, "released") empty = reset(4) next.ExpectedGeneration = empty.Generation if _, err := w.InitializeSandboxDeployment(t.Context(), next); err != nil { diff --git a/services/core/internal/store/sandbox_deployment_view_test.go b/services/core/internal/store/sandbox_deployment_view_test.go index 6414b5ceb..452633018 100644 --- a/services/core/internal/store/sandbox_deployment_view_test.go +++ b/services/core/internal/store/sandbox_deployment_view_test.go @@ -20,7 +20,7 @@ func TestSandboxDeploymentViewRecordsTemplateBuildAndSuspension(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) changes := deploymentExecution(t, w) id := uuid.NewString() diff --git a/services/core/internal/store/sandbox_deployment_worker_test.go b/services/core/internal/store/sandbox_deployment_worker_test.go index 039441dd4..70bdd2cb4 100644 --- a/services/core/internal/store/sandbox_deployment_worker_test.go +++ b/services/core/internal/store/sandbox_deployment_worker_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" @@ -50,7 +51,7 @@ func TestSandboxDeploymentWorkerActivatesWithoutRestart(t *testing.T) { if _, err := w.InitializeSandboxDeployment(t.Context(), sandbox.Selection{DeploymentSpec: store.SandboxDeploymentTestSpec("docker"), Provider: "docker"}); err != nil { t.Fatal(err) } - if _, err := s.CreateSession(t.Context(), uuid.NewString(), input); !errors.Is(err, deployment.ErrNodeUnavailable) { + if _, err := s.CreateSession(t.Context(), uuid.NewString(), input); !errors.Is(err, placement.ErrNodeUnavailable) { t.Fatal("zero-node deployment admitted Session", err) } token, err := store.EnrollmentTestToken(deployments.CreateEnrollment(t.Context(), deployment.Capacity{MaxActive: 4, MaxRetained: 16})) diff --git a/services/core/internal/store/sandbox_generations_test.go b/services/core/internal/store/sandbox_generations_test.go index 75c3d6e2c..42dfb6a7d 100644 --- a/services/core/internal/store/sandbox_generations_test.go +++ b/services/core/internal/store/sandbox_generations_test.go @@ -57,13 +57,13 @@ func TestE2BGenerationsRetainOwnershipAndUseCurrentCredential(t *testing.T) { if err != nil || len(generations) != 1 || generations[0].Generation != 1 || generations[0].Configuration.(*e2b.DeploymentConfiguration).APIURL != "https://api.e2b.app" { t.Fatal(generations, err) } - if _, err = w.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err = deploymentExecution(t, w).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err = w.SettleRuntimeCreation(t.Context(), owner); err != nil { + if _, err = deploymentExecution(t, w).SettleCreation(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err = w.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err = deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } if err = changes.CollectGenerations(t.Context()); err != nil { @@ -73,7 +73,7 @@ func TestE2BGenerationsRetainOwnershipAndUseCurrentCredential(t *testing.T) { if err != nil || len(generations) != 0 { t.Fatal(generations, err) } - historical, err := s.GetRuntimeAllocation(t.Context(), tenant, owner.EnvironmentID) + historical, err := deploymentStore(s).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: owner.EnvironmentID}) if err != nil || historical.DeploymentGeneration != 1 { t.Fatal("historical generation erased", err) } @@ -231,13 +231,13 @@ func TestGenerationDowngradeRefusesOldAllocation(t *testing.T) { if _, err = migrations.Up(t.Context()); err != nil { t.Fatal(err) } - if _, err = w.RequestRuntimeCleanup(t.Context(), owner); err != nil { + if _, err = deploymentExecution(t, w).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err = w.SettleRuntimeCreation(t.Context(), owner); err != nil { + if _, err = deploymentExecution(t, w).SettleCreation(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err = w.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err = deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } if _, err = migrations.DownTo(t.Context(), 80); err != nil { 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 17837417d..0e667702c 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 @@ -22,7 +22,7 @@ func TestSandboxNodeRoutesAuthenticateBeforeDeploymentState(t *testing.T) { if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, s, db), withCoreKeys(admin)) + handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), storeExecution(t, s), managedSandboxes(t, db), withCoreKeys(admin)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/sandbox_reset_test.go b/services/core/internal/store/sandbox_reset_test.go index 9c4648e4b..44a54f2bc 100644 --- a/services/core/internal/store/sandbox_reset_test.go +++ b/services/core/internal/store/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/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" @@ -70,7 +71,7 @@ func TestSandboxResetAutoUsesStartedWorkAndLockedRecheck(t *testing.T) { t.Run(kind, func(t *testing.T) { s, w, installation := managedArchiveFixture(t) tenant, session := managedArchiveSession(t, s, managerSessionInput(uuid.NewString())) - var owner RuntimeAllocation + var owner deployment.Allocation if kind != "pending" { owner = archiveAllocation(t, w, tenant, session, installation) } @@ -209,7 +210,7 @@ func TestSandboxResetAutoRechecksTurnStartedAfterListing(t *testing.T) { t.Fatal("listed idle candidate cut a new Turn", err) } before := adminMutationSnapshot(t, s, "sessions", "environments", "environment_input_reservations", "turns", "runtime_placements") - if _, err := s.CreateSession(t.Context(), tenant, managerSessionInput(uuid.NewString())); !errors.Is(err, ErrSandboxResetAdmission) { + if _, err := s.CreateSession(t.Context(), tenant, managerSessionInput(uuid.NewString())); !errors.Is(err, placement.ErrResetAdmission) { t.Fatal("new hosted admission during reset", err) } after := adminMutationSnapshot(t, s, "sessions", "environments", "environment_input_reservations", "turns", "runtime_placements") diff --git a/services/core/internal/store/sandbox_specification_store_test.go b/services/core/internal/store/sandbox_specification_store_test.go index 089ed0d8e..06464c77b 100644 --- a/services/core/internal/store/sandbox_specification_store_test.go +++ b/services/core/internal/store/sandbox_specification_store_test.go @@ -8,6 +8,7 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" @@ -23,7 +24,7 @@ func webSpecificationFixture(t *testing.T, provider string) (*Store, *Store, dep if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) w := executionWriter(t, s) id := uuid.NewString() changes := deploymentExecution(t, w) @@ -49,7 +50,7 @@ func specificationNode(t *testing.T, s *Store, view deployment.View) deployment. t.Fatal(err) } input := deployment.Enrollment{NodeID: uuid.NewString(), Name: "specification fixture", Credential: strings.Repeat("n", 64), Provider: view.Provider, - BackendFingerprint: strings.Repeat("b", 64), DeploymentGeneration: view.Generation, SpecificationDigest: view.SpecificationDigest, CoreURL: s.publicURL} + BackendFingerprint: strings.Repeat("b", 64), DeploymentGeneration: view.Generation, SpecificationDigest: view.SpecificationDigest, CoreURL: s.placement.PublicURL()} if _, err := nodes.Enroll(t.Context(), token, input); err != nil { t.Fatal(err) } @@ -135,14 +136,14 @@ func TestSandboxSpecificationChangesPreserveEveryRetainedResource(t *testing.T) if err != nil { t.Fatal(err) } - var owner RuntimeAllocation + var owner deployment.Allocation if state != "pending" { - owner, err = w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, view.InstallationID, runtimedevice.HashCredential(uuid.NewString())) + owner, err = deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, view.InstallationID, runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } if state != "creating" { - owner, err = w.ObserveRuntimeRunning(t.Context(), owner) + owner, err = deploymentExecution(t, w).ObserveRunning(t.Context(), owner) if err != nil { t.Fatal(err) } @@ -151,7 +152,7 @@ func TestSandboxSpecificationChangesPreserveEveryRetainedResource(t *testing.T) case "stopped": // A stopped native instance retains its allocation; the Store only // records that running compute could not be confirmed. - if err := w.RecordRuntimeObservation(t.Context(), owner, "compute_unconfirmed"); err != nil { + if err := deploymentExecution(t, w).RecordObservation(t.Context(), owner, "compute_unconfirmed"); err != nil { t.Fatal(err) } case "snapshot": @@ -161,7 +162,7 @@ func TestSandboxSpecificationChangesPreserveEveryRetainedResource(t *testing.T) t.Fatal(err) } case "cleanup_pending": - owner, err = w.RequestRuntimeCleanup(t.Context(), owner) + owner, err = deploymentExecution(t, w).RequestCleanup(t.Context(), owner) if err != nil { t.Fatal(err) } @@ -211,15 +212,15 @@ func TestSandboxSpecificationChangesPreserveEveryRetainedResource(t *testing.T) if !bytes.Equal(allocationBefore, allocationAfter) { t.Fatal("online change mutated or deleted retained ownership") } - owner, err = w.RequestRuntimeCleanup(t.Context(), owner) + owner, err = deploymentExecution(t, w).RequestCleanup(t.Context(), owner) if err != nil { t.Fatal(err) } - owner, err = w.SettleRuntimeCreation(t.Context(), owner) + owner, err = deploymentExecution(t, w).SettleCreation(t.Context(), owner) if err != nil { t.Fatal(err) } - if _, err := w.ReleaseRuntimeAllocation(t.Context(), owner); err != nil { + if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } } else if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { @@ -255,7 +256,7 @@ func TestSandboxSpecificationAllocationRaceWithMaintenance(t *testing.T) { } type result struct { session sessions.Session - owner RuntimeAllocation + owner deployment.Allocation err error } start := make(chan struct{}) @@ -264,7 +265,7 @@ func TestSandboxSpecificationAllocationRaceWithMaintenance(t *testing.T) { for _, session := range created { go func() { <-start - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, view.InstallationID, runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, view.InstallationID, runtimedevice.HashCredential(uuid.NewString())) results <- result{session, owner, err} }() } @@ -282,15 +283,15 @@ func TestSandboxSpecificationAllocationRaceWithMaintenance(t *testing.T) { result := <-results if result.err == nil { allocated++ - retry, err := w.ReserveRuntimeAllocation(t.Context(), tenant, result.session.Environment.ID, view.InstallationID, runtimedevice.HashCredential(uuid.NewString())) + retry, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: result.session.Environment.ID}, view.InstallationID, runtimedevice.HashCredential(uuid.NewString())) if err != nil || retry.ID != result.owner.ID || !retry.Replayed { t.Fatal("maintenance changed an admitted allocation retry", err) } } else { - if !errors.Is(result.err, ErrSandboxResetAdmission) { + if !errors.Is(result.err, placement.ErrResetAdmission) { t.Fatal("allocation race failed outside admission", result.err) } - if _, err := w.ReserveRuntimeAllocation(t.Context(), tenant, result.session.Environment.ID, view.InstallationID, runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, ErrSandboxResetAdmission) { + if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: result.session.Environment.ID}, view.InstallationID, runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, placement.ErrResetAdmission) { t.Fatal("fresh allocation passed committed maintenance", err) } } @@ -310,21 +311,21 @@ func TestSandboxSpecificationAllocationRaceWithMaintenance(t *testing.T) { // receives no new sandboxes until it is re-added. func TestNodeBoundToAnotherPublicURLGetsNoNewSandboxes(t *testing.T) { s, _, view, _ := webSpecificationFixture(t, "docker") - s.SetPublicURL("https://old.example") + s.SetPlacement(placementRules(t, "https://old.example")) node := specificationNode(t, s, view) nodes, err := deploymentService(t, s).ListNodes(t.Context()) if err != nil || len(nodes) != 1 || nodes[0].ID != node.NodeID || nodes[0].CoreURL != "https://old.example" { t.Fatal("enrollment did not record the node's address", nodes, err) } - s.SetPublicURL("https://new.example") - if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, deployment.ErrNodeUnavailable) { + s.SetPlacement(placementRules(t, "https://new.example")) + if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, placement.ErrNodeUnavailable) { t.Fatal("placed a new sandbox on a node bound to the old address", err) } bindings, err := deploymentService(t, s).AddressBindings(t.Context()) if err != nil || bindings.Nodes != 1 || bindings.NodesOnOtherAddress != 1 || bindings.HostedSandboxes != 0 { t.Fatal(bindings, err) } - s.SetPublicURL("https://old.example") + s.SetPlacement(placementRules(t, "https://old.example")) if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); err != nil { t.Fatal("node on the current address rejected placement", err) } @@ -334,14 +335,14 @@ func TestNodeBoundToAnotherPublicURLGetsNoNewSandboxes(t *testing.T) { // Session, and its configuration stays readable for cleanup. func TestE2BAdmitsNothingWhileThePublicURLIsLoopback(t *testing.T) { s, _, _, _ := webSpecificationFixture(t, "e2b") - s.SetPublicURL("http://127.0.0.1:8091") - if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, deployment.ErrPublicURLUnreachable) { + s.SetPlacement(placementRules(t, "http://127.0.0.1:8091")) + if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); !errors.Is(err, placement.ErrPublicURLUnreachable) { t.Fatal("admitted an E2B Session that could not reach Core", err) } if setup, err := deploymentService(t, s).Setup(t.Context()); err != nil || setup.Provider != "e2b" { t.Fatal("the saved E2B selection became unreadable", err) } - s.SetPublicURL("https://core.example") + s.SetPlacement(placementRules(t, "https://core.example")) if _, err := s.CreateSession(t.Context(), uuid.NewString(), managerSessionInput(uuid.NewString())); err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_creation_identity_test.go b/services/core/internal/store/session_creation_identity_test.go index 77983a3ff..b87bc54d1 100644 --- a/services/core/internal/store/session_creation_identity_test.go +++ b/services/core/internal/store/session_creation_identity_test.go @@ -134,7 +134,7 @@ func TestProviderKeyEntersRetryHashesOnlyAsKeyedFingerprint(t *testing.T) { if err != nil { t.Fatal(err) } - session, err := NewWithCredentialCipher(pool, cipher).CreateSession(t.Context(), uuid.NewString(), input) + session, err := withPlacement(t, NewWithCredentialCipher(pool, cipher)).CreateSession(t.Context(), uuid.NewString(), input) if err != nil { t.Fatal(err) } @@ -148,7 +148,7 @@ func TestProviderKeyEntersRetryHashesOnlyAsKeyedFingerprint(t *testing.T) { cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{71}, 32)) unreadable := input unreadable.IdempotencyKey, unreadable.CreationRequest = "unreadable", json.RawMessage(`{"x_agents_core":{"model_provider":"hash-key-canary"}}`) - if _, err := NewWithCredentialCipher(pool, cipher).CreateSession(t.Context(), uuid.NewString(), unreadable); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := withPlacement(t, NewWithCredentialCipher(pool, cipher)).CreateSession(t.Context(), uuid.NewString(), unreadable); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal("unreadable provider intent was hashed", err) } } diff --git a/services/core/internal/store/session_diagnostics_test.go b/services/core/internal/store/session_diagnostics_test.go index 582779d87..5de8dba1f 100644 --- a/services/core/internal/store/session_diagnostics_test.go +++ b/services/core/internal/store/session_diagnostics_test.go @@ -10,6 +10,7 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" @@ -206,7 +207,7 @@ func TestDiagnosticProvisioningDetailAtomicAndPrivate(t *testing.T) { t.Fatal(err) } writer := executionWriter(t, s) - owner, err := writer.ReserveRuntimeAllocation(t.Context(), tenant, session.Environment.ID, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, writer).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_execution_configuration_test.go b/services/core/internal/store/session_execution_configuration_test.go index e44de2a52..e15ef2f49 100644 --- a/services/core/internal/store/session_execution_configuration_test.go +++ b/services/core/internal/store/session_execution_configuration_test.go @@ -11,6 +11,7 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -36,7 +37,7 @@ func TestSessionExecutionConfigurationFrozenAcrossCreationPathsAndRetry(t *testi if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) for _, stream := range []bool{false, true} { for _, source := range []string{"session", "agent", "deployment"} { t.Run(source+map[bool]string{false: "/ordinary", true: "/stream"}[stream], func(t *testing.T) { @@ -164,7 +165,7 @@ func TestSessionExecutionConfigurationRollbackAndValidation(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) tenant := uuid.NewString() for _, kind := range []string{"model", "harness", "source", "provider", "provider_mismatch", "post_projection_failure"} { input := executionProjectionInput("session") @@ -260,15 +261,15 @@ func TestSessionExecutionConfigurationSurvivesSuspendResume(t *testing.T) { if err != nil { t.Fatal(err) } - owner, err := w.ReserveRuntimeAllocation(t.Context(), tenant, environment.ID, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) + owner, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}, uuid.NewString(), runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - owner, err = w.ObserveRuntimeRunning(t.Context(), owner) + owner, err = deploymentExecution(t, w).ObserveRunning(t.Context(), owner) if err != nil { t.Fatal(err) } - owner, err = w.SetRuntimeCompute(t.Context(), owner, "running", json.RawMessage(`{"instance":"original"}`), nil, 0) + owner, err = deploymentExecution(t, w).SetCompute(t.Context(), owner, "running", json.RawMessage(`{"instance":"original"}`), nil, 0) if err != nil { t.Fatal(err) } @@ -280,7 +281,7 @@ func TestSessionExecutionConfigurationSurvivesSuspendResume(t *testing.T) { retained = nil } owner = runtimeSuspensionStep(t, w, owner, phase, retained) - before, err := w.RuntimeActivity(t.Context(), owner) + before, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil { t.Fatal(err) } @@ -288,7 +289,7 @@ func TestSessionExecutionConfigurationSurvivesSuspendResume(t *testing.T) { if err != nil || !reflect.DeepEqual(got, frozen) { t.Fatal("runtime transition changed execution projection", phase, err) } - after, err := w.RuntimeActivity(t.Context(), owner) + after, err := deploymentStore(w).Activity(t.Context(), owner.ID) if err != nil || !before.LastActivity.Equal(after.LastActivity) || before.WakeRequested != after.WakeRequested || before.Busy != after.Busy || before.HasCompletedTurn != after.HasCompletedTurn { t.Fatal("configuration read changed runtime activity", err) } diff --git a/services/core/internal/store/session_initial_input.go b/services/core/internal/store/session_initial_input.go index e08ef4486..15c9642ed 100644 --- a/services/core/internal/store/session_initial_input.go +++ b/services/core/internal/store/session_initial_input.go @@ -3,6 +3,7 @@ package store import ( "context" "encoding/json" + "errors" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" @@ -11,8 +12,10 @@ import ( "github.com/jackc/pgx/v5/pgtype" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/placementpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/writeaudit" @@ -40,18 +43,29 @@ func (s *Store) createSessionResources(ctx context.Context, tenant string, param return err } if row.ID == params.ID { - var placement struct { + var configuration struct { Environment struct { Type string `json:"type"` } `json:"environment"` } - if err := json.Unmarshal(row.Configuration, &placement); err != nil { + if err := json.Unmarshal(row.Configuration, &configuration); err != nil { return err } - if placement.Environment.Type == "openai_hosted" { - if err := checkRuntimeDeploymentAdmission(ctx, q, ""); err != nil { + // Hosted admission and placement share the deployment lock with + // deployment changes, restore and node removal. + var hosted *placement.Deployment + if configuration.Environment.Type == "openai_hosted" { + if s.placement == nil { + return errors.New("hosted Session creation requires placement rules") + } + d, err := placementpg.LockDeployment(ctx, q) + if err != nil { return err } + if err := s.placement.CheckAdmission(d, ""); err != nil { + return err + } + hosted = &d } setup, err = s.freezeEnvironmentSkills(ctx, q, tenant, setup) if err != nil { @@ -107,10 +121,21 @@ func (s *Store) createSessionResources(ctx context.Context, tenant string, param return err } } - if placement.Environment.Type == "openai_hosted" { - if err := reserveRuntimePlacement(ctx, q, row.ID, s.publicURL); err != nil { + if hosted != nil { + nodes, err := placementpg.LoadNodes(ctx, q) + if err != nil { return err } + // A Session creation retry never reaches the placement. + chosen, err := s.placement.DecidePlacement(*hosted, nodes) + if err != nil { + return err + } + if chosen != nil { + if err := placementpg.ReservePlacement(ctx, q, row.ID, *chosen); err != nil { + return err + } + } } } environment, err = sessionEnvironmentSnapshot(ctx, q, row) 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 c50a0338d..e0b763697 100644 --- a/services/core/internal/store/session_model_execution_http_test.go +++ b/services/core/internal/store/session_model_execution_http_test.go @@ -17,9 +17,10 @@ func TestModelExecutionHTTPWriteOnlyAndStrictAdmission(t *testing.T) { _, pool := store.NewTestStore(t) cipher, _ := credentialcrypto.New(bytes.Repeat([]byte{6}, 32)) st, db := store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} + st.SetPlacement(fixtureRules(t, db)) tenant, token := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "catalog-test", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}}) - handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, st, db)) + handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, db)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_model_execution_test.go b/services/core/internal/store/session_model_execution_test.go index ceaca5a57..e5eb6ec92 100644 --- a/services/core/internal/store/session_model_execution_test.go +++ b/services/core/internal/store/session_model_execution_test.go @@ -18,7 +18,7 @@ func TestSessionModelExecutionEncryptedAndBound(t *testing.T) { if err != nil { t.Fatal(err) } - st := NewWithCredentialCipher(pool, cipher) + st := withPlacement(t, NewWithCredentialCipher(pool, cipher)) ctx := t.Context() tenant := uuid.NewString() input := sessions.CreateSession{Creator: FixtureCreator(), Engine: "mcode", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"model":"actual-model"},"environment":{"type":"openai_hosted"}}`), ModelProvider: &v1.ModelProviderInput{Protocol: "anthropic", BaseURL: "https://example.com", APIKey: "private-model-canary", ContextWindow: 100000, MaxOutputTokens: 8000}} @@ -41,7 +41,7 @@ func TestSessionModelExecutionEncryptedAndBound(t *testing.T) { if _, err := st.CreateSession(ctx, tenant, input); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal("changed credentials accepted", err) } - restarted := NewWithCredentialCipher(pool, cipher) + restarted := withPlacement(t, NewWithCredentialCipher(pool, cipher)) provider, err := restarted.SessionModelExecution(ctx, tenant, session.ID) if err != nil || provider.APIKey != "private-model-canary" { t.Fatal("restart lost model credential", err) @@ -56,7 +56,7 @@ func TestSessionModelExecutionEncryptedAndBound(t *testing.T) { t.Fatal("missing cipher succeeded") } input.IdempotencyKey = uuid.NewString() - if _, err := New(pool).CreateSession(ctx, tenant, input); !errors.Is(err, credentialcrypto.ErrUnavailable) { + if _, err := withPlacement(t, New(pool)).CreateSession(ctx, tenant, input); !errors.Is(err, credentialcrypto.ErrUnavailable) { t.Fatal("unencrypted create", err) } var count int diff --git a/services/core/internal/store/session_model_options_test.go b/services/core/internal/store/session_model_options_test.go index c7e2bc261..19e76c80c 100644 --- a/services/core/internal/store/session_model_options_test.go +++ b/services/core/internal/store/session_model_options_test.go @@ -17,7 +17,7 @@ func TestSessionModelExecutionStoresOnlyProviderBundle(t *testing.T) { if err != nil { t.Fatal(err) } - st := NewWithCredentialCipher(pool, cipher) + st := withPlacement(t, NewWithCredentialCipher(pool, cipher)) ctx, tenant := t.Context(), uuid.NewString() provider := &v1.ModelProviderInput{Protocol: "responses", BaseURL: "https://example.com/v1", APIKey: "provider-key-canary"} input := sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: []byte(`{"agent":{"model":"actual-model"},"environment":{"type":"openai_hosted"}}`), ModelProvider: provider} diff --git a/services/core/internal/store/sessions.go b/services/core/internal/store/sessions.go index 03fc02b09..28b66381a 100644 --- a/services/core/internal/store/sessions.go +++ b/services/core/internal/store/sessions.go @@ -18,6 +18,7 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment/placement" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/metadata" @@ -38,11 +39,15 @@ type Store struct { writer pgunit.Transactor lease *pgunit.Lease credentialCipher *credentialcrypto.Cipher - // publicURL is OAC_PUBLIC_URL. Core derives every address it gives - // nodes, sandboxes and administrators from it. - publicURL string + // placement admits and places hosted Session creation. It is nil until + // SetPlacement, and hosted creation fails without it. + placement *placement.Rules } +// SetPlacement records the installation's placement rules, built once in the +// composition root. Call it once, before serving requests. +func (s *Store) SetPlacement(rules *placement.Rules) { s.placement = rules } + func New(pool *pgxpool.Pool) *Store { pooled := pgunit.NewPool(pool) return &Store{queries: sqlc.New(pool), pool: pool, pooled: pooled, writer: pooled} diff --git a/services/core/internal/store/sessions_test.go b/services/core/internal/store/sessions_test.go index 4cb335a57..09d4bdd62 100644 --- a/services/core/internal/store/sessions_test.go +++ b/services/core/internal/store/sessions_test.go @@ -19,7 +19,7 @@ import ( func testStore(t *testing.T) (*Store, *pgxpool.Pool) { t.Helper() pool := pgtest.Open(t) - return New(pool), pool + return withPlacement(t, New(pool)), pool } // sessionAdapter is the Session adapter on s's database with s's credential diff --git a/services/core/internal/store/skill_version_deletion_test.go b/services/core/internal/store/skill_version_deletion_test.go index 6f75c10d4..f98a573c6 100644 --- a/services/core/internal/store/skill_version_deletion_test.go +++ b/services/core/internal/store/skill_version_deletion_test.go @@ -21,7 +21,7 @@ func TestSoleSkillVersionDeletionKeepsFrozenSetup(t *testing.T) { if err != nil { t.Fatal(err) } - s := NewWithCredentialCipher(pool, cipher) + s := withPlacement(t, NewWithCredentialCipher(pool, cipher)) skillService := SkillService(t, pool, cipher) tenant := uuid.NewString() archive := skillArchive(t, "sole-version-frozen") diff --git a/services/core/internal/store/template_composition_public_test.go b/services/core/internal/store/template_composition_public_test.go index c4a697f2f..1585ea31b 100644 --- a/services/core/internal/store/template_composition_public_test.go +++ b/services/core/internal/store/template_composition_public_test.go @@ -33,6 +33,8 @@ func TestTemplateCompositionOfficialClientPostgres(t *testing.T) { } s, reopenedStore := store.NewWithCredentialCipher(pool, cipher), store.NewWithCredentialCipher(pool, cipher) db := fixtureDB{pool: pool, cipher: cipher} // built both Stores + s.SetPlacement(fixtureRules(t, db)) + reopenedStore.SetPlacement(fixtureRules(t, db)) tenant, foreignTenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "composition-owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, @@ -41,7 +43,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, db, auth, "codex", storeExecution(t, current), managedSandboxes(t, current, db), fixtureDeploymentProvider()) + h, err := publicHandler(t, current, db, auth, "codex", storeExecution(t, current), managedSandboxes(t, db), fixtureDeploymentProvider()) 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 fb4a777ee..edfe2a46e 100644 --- a/services/core/internal/store/template_null_selection_public_test.go +++ b/services/core/internal/store/template_null_selection_public_test.go @@ -35,6 +35,8 @@ func TestTemplateNullSelectionOfficialClientPostgres(t *testing.T) { } s, reopenedStore := store.NewWithCredentialCipher(pool, cipher), store.NewWithCredentialCipher(pool, cipher) db := fixtureDB{pool: pool, cipher: cipher} // built both Stores + s.SetPlacement(fixtureRules(t, db)) + reopenedStore.SetPlacement(fixtureRules(t, db)) tenant, foreignTenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{ {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "selection-owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, @@ -42,7 +44,7 @@ func TestTemplateNullSelectionOfficialClientPostgres(t *testing.T) { }) serve := func(current *store.Store) *httptest.Server { t.Helper() - h, err := publicHandler(t, current, db, auth, "codex", storeExecution(t, current), managedSandboxes(t, current, db), fixtureDeploymentProvider()) + h, err := publicHandler(t, current, db, auth, "codex", storeExecution(t, current), managedSandboxes(t, db), fixtureDeploymentProvider()) 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 4d17474f2..9c7d4a57b 100644 --- a/services/core/internal/store/unified_model_configuration_http_test.go +++ b/services/core/internal/store/unified_model_configuration_http_test.go @@ -21,7 +21,7 @@ func TestUnifiedModelConfigurationHTTP(t *testing.T) { if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, st, db), withCoreKeys(admin), withHarnesses([]string{"codex", "claude_sdk"})) + handler, err := publicHandler(t, st, db, auth, "codex", storeExecution(t, st), managedSandboxes(t, db), withCoreKeys(admin), withHarnesses([]string{"codex", "claude_sdk"})) if err != nil { t.Fatal(err) }