From 2eee6e4c28177c8eb37f826b46bdc3a2945900d6 Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 30 Sep 2026 21:52:05 +0000 Subject: [PATCH] Move Environments and devices into sessions and sessionpg Environment and device reads, the initialization list and the frozen setup and initial files join sessions.Reader on the pooled sessionpg.Store, which sessionpg.New(units, cipher) builds with the credential key. Frozen data that does not open or validate is an internal error, and a missing key is credentialcrypto.ErrUnavailable. Device creation, revocation, heartbeats and Runtime enrollment are sessions.Service use cases; the device registration rule and RuntimeTypeAgentDaemon move to sessions and runtimedevice. The EnvironmentExecution family joins sessions.ExecutionOperations on sessionpg.NewExecution(lease): initialization claim, completion and failure, connection replacement, observation and reconciliation, and Session device binding each run a Load, Decide and Apply transaction under the Session lock. sessionpg.LoadEnvironment, LoadSessionEnvironment and CreateSessionEnvironment are participants of other Session transactions. The Worker, runtime manager and lifecycle receive the Session operations and reader explicitly. The daemon gateway takes its credential, heartbeat and archived cancellation dependencies separately, so draining a receipt no longer depends on a type assertion on the heartbeat store, and MarkRuntimeOffline is deleted. api gains the required EnvironmentsReader, and cmd/device provisions devices through sessions.Service. The store Environment, initialization, connection recovery, local device and enrollment files are deleted. --- services/core/IMPLEMENTATION.md | 4 +- services/core/cmd/device/main.go | 11 +- .../core/cmd/server/daemon_bootstrap_test.go | 2 +- .../core/cmd/server/executor_connections.go | 6 +- services/core/cmd/server/http_routes_test.go | 2 +- services/core/cmd/server/main.go | 12 +- services/core/internal/api/dependencies.go | 3 +- .../core/internal/api/dependencies_test.go | 5 +- .../core/internal/api/environment_files.go | 2 +- .../internal/api/environment_files_create.go | 2 +- .../api/environment_files_create_test.go | 2 +- .../internal/api/environment_files_test.go | 2 +- .../internal/api/environment_installation.go | 2 +- services/core/internal/api/environments.go | 20 +- .../core/internal/api/environments_test.go | 2 +- services/core/internal/api/fakes_test.go | 20 +- .../session_diagnostics_public_compat_test.go | 2 +- .../archive_cancellation_cleanup_test.go | 6 +- .../internal/execution/device_authority.go | 7 +- .../execution/environment_connections.go | 13 - .../execution/environment_directory.go | 6 +- .../execution/environment_file_write.go | 4 +- .../core/internal/execution/message_input.go | 1 + .../internal/execution/prepared_dispatch.go | 2 +- .../internal/execution/runtime_compute.go | 4 +- .../execution/runtime_compute_wake.go | 4 +- .../internal/execution/runtime_connections.go | 16 +- .../execution/runtime_initialization.go | 18 +- .../internal/execution/runtime_lifecycle.go | 38 +- .../internal/execution/runtime_manager.go | 5 +- .../internal/execution/runtime_wake_hint.go | 2 +- .../sandbox_deployment_drain_test.go | 4 +- .../sandbox_deployment_setup_test.go | 12 +- .../sandbox_deployment_switch_test.go | 8 +- .../execution/sandbox_generations_test.go | 2 +- .../sandbox_provider_contract_test.go | 2 +- .../internal/execution/sandbox_reset_test.go | 24 +- .../execution/sandbox_snapshot_budget_test.go | 2 +- .../execution/session_fixture_test.go | 24 ++ services/core/internal/execution/worker.go | 4 +- .../core/internal/execution/worker_device.go | 9 +- .../postgres/sessionpg/artifacts_test.go | 2 +- .../persistence/postgres/sessionpg/devices.go | 225 +++++++++++ .../postgres/sessionpg/devices_test.go | 71 ++++ .../postgres/sessionpg/environment.go | 83 ++++ .../sessionpg/execution_environment.go | 164 ++++++++ .../sessionpg/executor_credentials.go | 35 ++ .../persistence/postgres/sessionpg/frozen.go | 78 ++++ .../postgres/sessionpg/frozen_test.go | 141 +++++++ .../persistence/postgres/sessionpg/session.go | 25 +- .../persistence/postgres/sessionpg/store.go | 16 +- services/core/internal/runtime/gateway.go | 18 +- .../core/internal/runtimedevice/credential.go | 4 + .../runtimeenrollment/connection_test.go | 2 +- .../internal/runtimeenrollment/enrollment.go | 3 +- .../runtimeenrollment/enrollment_test.go | 7 +- services/core/internal/runtimegateway/auth.go | 8 +- .../core/internal/runtimegateway/auth_test.go | 4 +- .../runtimegateway/bootstrap_url_test.go | 2 +- .../internal/runtimegateway/cancellation.go | 5 +- .../runtimegateway/cancellation_test.go | 7 +- .../core/internal/runtimegateway/handler.go | 6 +- .../internal/runtimegateway/handler_test.go | 2 +- .../mcp_bearer_live_linux_test.go | 2 +- .../core/internal/runtimegateway/session.go | 15 +- .../internal/runtimegateway/session_test.go | 4 - .../core/internal/runtimegateway/wire_test.go | 2 +- .../providers/configuration_flow_test.go | 2 +- .../core/internal/sessions/artifacts_test.go | 25 ++ services/core/internal/sessions/devices.go | 177 ++++++++ .../core/internal/sessions/devices_test.go | 194 +++++++++ .../core/internal/sessions/environment.go | 53 +++ .../internal/sessions/environment_test.go | 25 ++ services/core/internal/sessions/execution.go | 1 + .../sessions/execution_environment.go | 355 ++++++++++++++++ .../sessions/execution_environment_test.go | 380 ++++++++++++++++++ .../sessions/execution_functions_test.go | 12 +- .../internal/sessions/executor_credentials.go | 10 + services/core/internal/sessions/reader.go | 7 +- services/core/internal/sessions/service.go | 1 + .../internal/sessions/transaction_test.go | 17 + .../store/admin_session_archive_test.go | 4 +- .../store/archive_cancellation_test.go | 16 +- .../internal/store/claude_execution_test.go | 4 +- .../core/internal/store/claude_mcp_test.go | 2 +- .../creation_stream_settlement_public_test.go | 6 +- .../store/credential_matrix_http_test.go | 6 +- .../store/device_bootstrap_binding_test.go | 14 +- services/core/internal/store/devices.go | 135 ------- services/core/internal/store/devices_test.go | 45 ++- services/core/internal/store/dispatch_test.go | 12 +- .../environment_connection_events_test.go | 5 +- .../store/environment_connection_recovery.go | 50 --- .../environment_connection_recovery_test.go | 15 +- .../environment_connection_worker_test.go | 11 +- .../internal/store/environment_connections.go | 111 ----- .../store/environment_connections_test.go | 32 +- .../internal/store/environment_device_test.go | 6 +- .../store/environment_directory_test.go | 2 +- .../environment_executor_command_test.go | 8 +- .../store/environment_executor_credentials.go | 24 -- .../environment_executor_credentials_test.go | 12 +- .../store/environment_executor_management.go | 3 +- .../environment_executor_management_test.go | 14 +- .../internal/store/environment_file_writes.go | 2 +- .../store/environment_file_writes_test.go | 6 +- .../store/environment_initial_input_test.go | 8 +- .../store/environment_initialization.go | 116 ------ .../store/environment_initialization_test.go | 44 +- .../store/environment_input_activity_test.go | 22 +- .../store/environment_input_claim_test.go | 3 +- .../store/environment_installation_test.go | 4 +- .../store/environment_packages_test.go | 32 -- .../store/environment_plugins_test.go | 4 +- .../store/environment_retrieve_public_test.go | 4 +- .../store/environment_runtime_fixture_test.go | 12 +- .../core/internal/store/environment_setup.go | 39 -- .../internal/store/environment_setup_test.go | 4 +- .../environment_skill_references_test.go | 4 +- .../internal/store/environment_skills_test.go | 4 +- .../store/environment_steering_order_test.go | 3 +- .../internal/store/environment_work_test.go | 2 +- services/core/internal/store/environments.go | 60 --- .../core/internal/store/environments_test.go | 26 +- .../core/internal/store/execution_test.go | 18 +- .../store/executor_credential_target.go | 3 +- .../executor_principals_migration_test.go | 2 +- .../store/executor_principals_test.go | 10 +- .../store/function_execution_native_test.go | 2 +- .../internal/store/function_execution_test.go | 4 +- .../internal/store/function_worker_test.go | 4 +- ...sted_initialization_failure_public_test.go | 14 +- services/core/internal/store/initial_files.go | 19 - .../internal/store/initial_files_http_test.go | 2 +- .../core/internal/store/initial_files_test.go | 6 +- .../internal/store/list_cursor_public_test.go | 20 +- .../store/local_environment_devices.go | 38 -- .../store/local_environment_devices_test.go | 29 +- .../store/local_environment_worker_test.go | 6 +- .../store/message_input_helpers_test.go | 3 +- .../internal/store/native_environment_test.go | 2 +- .../store/native_public_execution_test.go | 4 +- .../internal/store/native_recovery_test.go | 11 +- .../internal/store/no_environment_test.go | 2 +- .../store/provider_operations_fixture_test.go | 1 + .../internal/store/public_execution_test.go | 2 +- .../store/public_handler_fixture_test.go | 11 +- .../internal/store/runtime_allocations.go | 11 +- .../store/runtime_allocations_test.go | 6 +- .../runtime_capabilities_pending_test.go | 6 +- .../store/runtime_compute_lifecycle_test.go | 10 +- .../runtime_configuration_cleanup_test.go | 4 +- .../internal/store/runtime_connection_test.go | 8 +- .../store/runtime_creation_settlement_test.go | 8 +- .../store/runtime_deployment_worker_test.go | 2 +- .../core/internal/store/runtime_enrollment.go | 91 ----- .../runtime_enrollment_connection_test.go | 10 +- .../internal/store/runtime_enrollment_test.go | 32 +- .../runtime_environment_terminal_test.go | 6 +- .../store/runtime_history_acceptance_test.go | 14 +- .../store/runtime_initialization_test.go | 21 +- .../store/runtime_input_admission_test.go | 2 +- .../internal/store/runtime_lifecycle_test.go | 22 +- .../runtime_node_lifecycle_fixture_test.go | 4 +- .../store/runtime_node_lifecycle_test.go | 4 +- .../store/runtime_observation_scan_test.go | 4 +- .../internal/store/runtime_pending_test.go | 4 +- .../core/internal/store/runtime_scan_test.go | 6 +- .../core/internal/store/runtime_suspension.go | 2 +- .../store/runtime_worker_recovery_test.go | 2 +- .../store/sandbox_deployment_switch_test.go | 2 +- .../sandbox_deployment_switch_worker_test.go | 2 +- .../store/sandbox_deployment_worker_test.go | 2 +- .../core/internal/store/sandbox_reset_test.go | 2 +- .../store/self_hosted_initial_public_test.go | 2 +- .../internal/store/session_artifacts_test.go | 6 +- .../session_deletion_lifecycle_public_test.go | 4 +- .../store/session_devices_fixture_test.go | 33 ++ .../store/session_diagnostics_test.go | 4 +- .../session_execution_configuration_test.go | 2 +- .../internal/store/session_initial_input.go | 8 +- .../store/session_reads_fixture_test.go | 43 ++ .../store/session_write_audit_test.go | 2 +- .../internal/store/sessions_fixture_test.go | 29 +- services/core/internal/store/sessions_test.go | 17 +- .../store/skill_version_deletion_test.go | 4 +- .../store/structured_output_dispatch_test.go | 2 +- .../internal/store/subagent_dispatch_test.go | 2 +- .../store/subagent_identities_test.go | 6 +- .../store/subagent_native_outputs_test.go | 8 +- .../internal/store/subagent_resources_test.go | 4 +- .../store/subagent_visibility_public_test.go | 4 +- .../store/template_composition_public_test.go | 44 +- .../template_null_selection_public_test.go | 68 ++-- services/core/internal/store/turns_test.go | 3 +- 195 files changed, 2877 insertions(+), 1360 deletions(-) delete mode 100644 services/core/internal/execution/environment_connections.go create mode 100644 services/core/internal/execution/session_fixture_test.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/devices.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/devices_test.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/execution_environment.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/executor_credentials.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/frozen.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/frozen_test.go create mode 100644 services/core/internal/sessions/devices.go create mode 100644 services/core/internal/sessions/devices_test.go create mode 100644 services/core/internal/sessions/environment_test.go create mode 100644 services/core/internal/sessions/execution_environment.go create mode 100644 services/core/internal/sessions/execution_environment_test.go delete mode 100644 services/core/internal/store/environment_connection_recovery.go delete mode 100644 services/core/internal/store/environment_initialization.go delete mode 100644 services/core/internal/store/environment_packages_test.go delete mode 100644 services/core/internal/store/environments.go delete mode 100644 services/core/internal/store/local_environment_devices.go delete mode 100644 services/core/internal/store/runtime_enrollment.go create mode 100644 services/core/internal/store/session_devices_fixture_test.go create mode 100644 services/core/internal/store/session_reads_fixture_test.go diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index b971bb139..5253735bf 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 two `sessionpg` participants, on their own transaction's queries: `LockSession`, which orders a write against the Session's Turns, and `BindSession`, which runs a Session procedure, such as `CreateEnvironmentDevice`, in that transaction. 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` and `LoadSessionEnvironment`, which read the tenant's Environment of a live Session by Environment or by 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`. `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`. @@ -34,7 +34,7 @@ Domain owners, each with its PostgreSQL adapter under `internal/persistence/post - `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`. -- `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) and Session Artifacts (their reads, deletion and the staging of a Turn's export). `cmd/server` builds one `sessionpg.Store` and wires the Service and the Store into `execution.Dispatcher.Sessions` and `SessionsReader` and into the api fields; the Worker stages Artifacts through the Service. Function calls and their application receipts and the Turn execution journal run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. +- `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) and executor credential authentication (`sessions.ExecutorCredentialReader`). `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. Function calls and their application receipts, the Turn execution journal, Environment initialization, connection observations and their reconciliation, and Session device binding run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. ## Request handling diff --git a/services/core/cmd/device/main.go b/services/core/cmd/device/main.go index df1086680..59c6830ce 100644 --- a/services/core/cmd/device/main.go +++ b/services/core/cmd/device/main.go @@ -17,8 +17,10 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/obs/log" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/databaseurl" + "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/runtimedevice" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) func main() { @@ -54,7 +56,12 @@ func run() error { return errors.New("invalid execution database configuration") } defer pool.Close() - s := store.New(pool) + // Device provisioning opens no frozen Session data, so it needs no + // credential key. + s, err := sessions.NewService(sessionpg.New(pgunit.NewPool(pool), nil)) + if err != nil { + return err + } if *revoke != "" { return s.RevokeDevice(ctx, *tenant, *revoke) } diff --git a/services/core/cmd/server/daemon_bootstrap_test.go b/services/core/cmd/server/daemon_bootstrap_test.go index 69880a323..372df38ae 100644 --- a/services/core/cmd/server/daemon_bootstrap_test.go +++ b/services/core/cmd/server/daemon_bootstrap_test.go @@ -18,7 +18,7 @@ type bootstrapCredentialStore struct { } func (s bootstrapCredentialStore) GetDeviceCredential(context.Context, string) (runtimedevice.Credential, bool, error) { - return runtimedevice.Credential{ID: "runtime", Type: runtimegateway.RuntimeTypeAgentDaemon, + return runtimedevice.Credential{ID: "runtime", Type: runtimedevice.RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential("synthetic-token"), RuntimeNodeID: s.nodeID, RuntimeAllocationID: s.allocationID}, true, nil } diff --git a/services/core/cmd/server/executor_connections.go b/services/core/cmd/server/executor_connections.go index 9c8138c3f..bd38a0dcc 100644 --- a/services/core/cmd/server/executor_connections.go +++ b/services/core/cmd/server/executor_connections.go @@ -5,16 +5,16 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeenrollment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) // executorConnections observes enrolled executors through the Runtime gateway. // registry is nil when this Core has no gateway; no executor is then connected. type executorConnections struct { - store *store.Store + sessions sessions.Reader registry *runtimegateway.Registry } func (c executorConnections) ExecutorConnected(ctx context.Context, environment, digest string) (bool, error) { - return runtimeenrollment.RuntimeConnected(ctx, c.store, c.registry, environment, digest) + return runtimeenrollment.RuntimeConnected(ctx, c.sessions, c.registry, environment, digest) } diff --git a/services/core/cmd/server/http_routes_test.go b/services/core/cmd/server/http_routes_test.go index 4c9b2503e..6c2bed018 100644 --- a/services/core/cmd/server/http_routes_test.go +++ b/services/core/cmd/server/http_routes_test.go @@ -103,7 +103,7 @@ func daemonComposition(t testing.TB) http.Handler { Subagents: struct{ api.Subagents }{}, Artifacts: struct{ api.Artifacts }{}, ArtifactsReader: struct{ api.ArtifactsReader }{}, - SessionAdmin: struct{ api.SessionAdmin }{}, Environments: struct{ api.Environments }{}, ExecutorConnections: struct{ api.ExecutorConnections }{}, + SessionAdmin: struct{ api.SessionAdmin }{}, Environments: struct{ api.Environments }{}, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, ExecutorConnections: struct{ api.ExecutorConnections }{}, Admin: struct{ api.Admin }{}, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, Metrics: struct{ api.Metrics }{}, RuntimeObservations: struct{ api.RuntimeObservations }{}, RuntimeHistory: struct{ api.RuntimeHistory }{}, Execution: &api.Execution{ diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index dad7ba5b5..8c343b916 100644 --- a/services/core/cmd/server/main.go +++ b/services/core/cmd/server/main.go @@ -164,7 +164,7 @@ func run() error { if err != nil { return err } - sessionStore := sessionpg.New(units) + sessionStore := sessionpg.New(units, credentialKey) sessionService, err := sessions.NewService(sessionStore) if err != nil { return err @@ -248,7 +248,7 @@ func run() error { if err := api.ValidateCredentialSeparation(ctx, keyAdmin, projectStore); err != nil { return err } - historyResolver, err := historystoreresolver.NewResolver(executionStore) + historyResolver, err := historystoreresolver.NewResolver(sessionStore) if err != nil { return err } @@ -265,7 +265,7 @@ func run() error { if err != nil { return err } - daemonHandler, registry, err = runtime.NewGateway(executionStore, executorURL) + daemonHandler, registry, err = runtime.NewGateway(sessionStore, sessionService, executionStore, executorURL) if err != nil { return err } @@ -392,7 +392,7 @@ func run() error { Artifacts: sessionService, ArtifactsReader: sessionStore, SessionAdmin: executionStore, - Environments: executionStore, ExecutorConnections: executorConnections{store: executionStore, registry: registry}, + Environments: executionStore, EnvironmentsReader: sessionStore, ExecutorConnections: executorConnections{sessions: sessionStore, registry: registry}, Admin: executionStore, AdminAudit: auditStore, WriteAudit: auditStore, Metrics: metrics, RuntimeObservations: observationService, RuntimeHistory: historyService, } @@ -421,8 +421,8 @@ func run() error { } if daemonHandler != nil { routes := daemonRoutes{gateway: daemonHandler, - enrollment: runtimeenrollment.EnrollmentHandler(executionStore), - connection: runtimeenrollment.ConnectionHandler(executionStore, registry)} + enrollment: runtimeenrollment.EnrollmentHandler(sessionService), + connection: runtimeenrollment.ConnectionHandler(sessionStore, registry)} if managedNodes != nil { routes.nodeConnect = managedNodes.hub } diff --git a/services/core/internal/api/dependencies.go b/services/core/internal/api/dependencies.go index f5250ee05..d1e25c875 100644 --- a/services/core/internal/api/dependencies.go +++ b/services/core/internal/api/dependencies.go @@ -52,6 +52,7 @@ type Dependencies struct { ArtifactsReader ArtifactsReader SessionAdmin SessionAdmin Environments Environments + EnvironmentsReader EnvironmentsReader ExecutorConnections ExecutorConnections Admin Admin AdminAudit AdminAudit @@ -139,7 +140,7 @@ func (d Dependencies) validate() error { field{"Subagents", d.Subagents}, field{"Artifacts", d.Artifacts}, field{"ArtifactsReader", d.ArtifactsReader}, - field{"SessionAdmin", d.SessionAdmin}, field{"Environments", d.Environments}, + field{"SessionAdmin", d.SessionAdmin}, field{"Environments", d.Environments}, field{"EnvironmentsReader", d.EnvironmentsReader}, field{"ExecutorConnections", d.ExecutorConnections}, field{"Admin", d.Admin}, field{"AdminAudit", d.AdminAudit}, field{"WriteAudit", d.WriteAudit}, field{"Metrics", d.Metrics}, field{"RuntimeObservations", d.RuntimeObservations}, field{"RuntimeHistory", d.RuntimeHistory}, ); err != nil { diff --git a/services/core/internal/api/dependencies_test.go b/services/core/internal/api/dependencies_test.go index 246d1f931..41722853c 100644 --- a/services/core/internal/api/dependencies_test.go +++ b/services/core/internal/api/dependencies_test.go @@ -38,6 +38,7 @@ type testFakes struct { artifactsReader *fakeArtifactsReader sessionAdmin *fakeSessionAdmin environments *fakeEnvironments + environmentsReader *fakeEnvironmentsReader executorConnections *fakeExecutorConnections admin *fakeAdmin adminAudit *fakeAdminAudit @@ -82,7 +83,7 @@ func testDependencies(t testing.TB) (Dependencies, *testFakes) { subagents: &fakeSubagents{t: t}, artifacts: &fakeArtifacts{t: t}, artifactsReader: &fakeArtifactsReader{t: t}, - sessionAdmin: &fakeSessionAdmin{t: t}, environments: &fakeEnvironments{t: t}, executorConnections: &fakeExecutorConnections{t: t}, + sessionAdmin: &fakeSessionAdmin{t: t}, environments: &fakeEnvironments{t: t}, environmentsReader: &fakeEnvironmentsReader{t: t}, executorConnections: &fakeExecutorConnections{t: t}, admin: &fakeAdmin{t: t}, adminAudit: &fakeAdminAudit{t: t}, writeAudit: &fakeWriteAudit{t: t}, metrics: &fakeMetrics{t: t}, runtimeObservations: &fakeRuntimeObservations{t: t}, runtimeHistory: &fakeRuntimeHistory{t: t}, installationBindings: &fakeInstallationBindings{t: t}, sessionAdmission: &fakeSessionAdmission{t: t}, @@ -110,7 +111,7 @@ func testDependencies(t testing.TB) (Dependencies, *testFakes) { Artifacts: f.artifacts, ArtifactsReader: f.artifactsReader, SessionAdmin: f.sessionAdmin, - Environments: f.environments, ExecutorConnections: f.executorConnections, Admin: f.admin, AdminAudit: f.adminAudit, WriteAudit: f.writeAudit, + Environments: f.environments, EnvironmentsReader: f.environmentsReader, ExecutorConnections: f.executorConnections, Admin: f.admin, AdminAudit: f.adminAudit, WriteAudit: f.writeAudit, Metrics: f.metrics, RuntimeObservations: f.runtimeObservations, RuntimeHistory: f.runtimeHistory, }, f } diff --git a/services/core/internal/api/environment_files.go b/services/core/internal/api/environment_files.go index fbab274bf..35440fa9a 100644 --- a/services/core/internal/api/environment_files.go +++ b/services/core/internal/api/environment_files.go @@ -38,7 +38,7 @@ type EnvironmentWorkspaces interface { // @Failure 400,401,404,500,503 {object} v1.ErrorResponse // @Router /agents/environments/{environment_id}/files [get] func (h *Handler) listEnvironmentFiles(w http.ResponseWriter, r *http.Request) { - environment, err := h.Environments.GetEnvironment(r.Context(), tenantID(r), chi.URLParam(r, "environment_id")) + environment, err := h.EnvironmentsReader.GetEnvironment(r.Context(), tenantID(r), chi.URLParam(r, "environment_id")) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/environment_files_create.go b/services/core/internal/api/environment_files_create.go index 0bc6aaedd..dfb4e901a 100644 --- a/services/core/internal/api/environment_files_create.go +++ b/services/core/internal/api/environment_files_create.go @@ -40,7 +40,7 @@ func (h *Handler) createEnvironmentFile(w http.ResponseWriter, r *http.Request) if !ok { return } - environment, err := h.Environments.GetEnvironment(r.Context(), tenantID(r), chi.URLParam(r, "environment_id")) + environment, err := h.EnvironmentsReader.GetEnvironment(r.Context(), tenantID(r), chi.URLParam(r, "environment_id")) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/environment_files_create_test.go b/services/core/internal/api/environment_files_create_test.go index 5129de312..ebeebd16a 100644 --- a/services/core/internal/api/environment_files_create_test.go +++ b/services/core/internal/api/environment_files_create_test.go @@ -47,7 +47,7 @@ func environmentFileCreateHandler(t *testing.T, configure ...func(*Dependencies, APIKey{OrganizationID: "org", ProjectID: "project", SubjectKind: "user", SubjectID: "caller", TokenSHA256: runtimedevice.HashCredential("files-key"), TenantID: f.environment.TenantID}, APIKey{OrganizationID: "org", ProjectID: "other", SubjectKind: "user", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential("other-key"), TenantID: uuid.NewString()}, ).ResolveAPIKey - fakes.environments.getEnvironment = f.GetEnvironment + fakes.environmentsReader.getEnvironment = f.GetEnvironment deps.Execution = fakes.execution() fakes.workspaces.readEnvironmentDirectory, fakes.workspaces.writeEnvironmentFile = f.ReadEnvironmentDirectory, f.WriteEnvironmentFile for _, c := range configure { diff --git a/services/core/internal/api/environment_files_test.go b/services/core/internal/api/environment_files_test.go index 4ad1bffa0..4ba1ebad1 100644 --- a/services/core/internal/api/environment_files_test.go +++ b/services/core/internal/api/environment_files_test.go @@ -79,7 +79,7 @@ func environmentFilesHandler(t *testing.T, enabled bool, configure ...func(*Depe deps, fakes := testDependencies(t) deps.Engine = "fake_alpha" fakes.projectsReader.resolveAPIKey = projectKeys(t, keys...).ResolveAPIKey - fakes.environments.getEnvironment = f.GetEnvironment + fakes.environmentsReader.getEnvironment = f.GetEnvironment if enabled { deps.Execution = fakes.execution() fakes.workspaces.readEnvironmentDirectory = f.ReadEnvironmentDirectory diff --git a/services/core/internal/api/environment_installation.go b/services/core/internal/api/environment_installation.go index 121e23f6f..de20c3322 100644 --- a/services/core/internal/api/environment_installation.go +++ b/services/core/internal/api/environment_installation.go @@ -113,7 +113,7 @@ func (h *Handler) prepareNativeInstallation(w http.ResponseWriter, r *http.Reque if !ok { return } - environment, err := h.Environments.GetEnvironment(r.Context(), claim.Principal.TenantID, claim.Environment) + environment, err := h.EnvironmentsReader.GetEnvironment(r.Context(), claim.Principal.TenantID, claim.Environment) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/environments.go b/services/core/internal/api/environments.go index a44b7be0b..d212919c9 100644 --- a/services/core/internal/api/environments.go +++ b/services/core/internal/api/environments.go @@ -12,12 +12,16 @@ import ( "github.com/go-chi/chi/v5" ) -// Environments reads Environments, grants and claims native installations, and -// manages the executor credentials of a Project's self_hosted Environments. The -// Project's principal is an executor credential's execution principal; the -// Core key that authorizes the request is not. -type Environments interface { +// EnvironmentsReader reads Environments. +type EnvironmentsReader interface { GetEnvironment(context.Context, string, string) (sessions.Environment, error) +} + +// Environments grants and claims native installations, and manages the +// executor credentials of a Project's self_hosted Environments. The Project's +// principal is an executor credential's execution principal; the Core key that +// authorizes the request is not. +type Environments interface { AuthorizeEnvironmentInstallation(context.Context, identity.Principal, string, string) (string, int64, error) ValidateEnvironmentInstallation(context.Context, string, string) (sessions.InstallationAuthorization, error) ClaimEnvironmentInstallation(context.Context, string, string, string) error @@ -37,14 +41,14 @@ type Environments interface { // @Failure 400,401,404,500 {object} v1.ErrorResponse // @Router /agents/environments/{environment_id} [get] func (h *Handler) getEnvironment(w http.ResponseWriter, r *http.Request) { - environment, err := h.Environments.GetEnvironment(r.Context(), tenantID(r), chi.URLParam(r, "environment_id")) + environment, err := h.EnvironmentsReader.GetEnvironment(r.Context(), tenantID(r), chi.URLParam(r, "environment_id")) if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } response, err := environmentResponse(environment) if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } writeJSON(w, http.StatusOK, response) diff --git a/services/core/internal/api/environments_test.go b/services/core/internal/api/environments_test.go index 3104677d5..f7a8db708 100644 --- a/services/core/internal/api/environments_test.go +++ b/services/core/internal/api/environments_test.go @@ -40,7 +40,7 @@ func environmentResourceHandler(t *testing.T) (http.Handler, *environmentResourc OrganizationID: "resource-org", ProjectID: "resource-project", SubjectKind: "user", SubjectID: "resource-reader", TokenSHA256: runtimedevice.HashCredential("resource-key"), TenantID: f.environment.TenantID, }).ResolveAPIKey - fakes.environments.getEnvironment = f.GetEnvironment + fakes.environmentsReader.getEnvironment = f.GetEnvironment return newTestHandler(t, deps), f } diff --git a/services/core/internal/api/fakes_test.go b/services/core/internal/api/fakes_test.go index 2c6c71bcf..ac64473bd 100644 --- a/services/core/internal/api/fakes_test.go +++ b/services/core/internal/api/fakes_test.go @@ -393,9 +393,20 @@ func (f *fakeEnvironmentWorkspaces) WriteEnvironmentFile(a0 context.Context, a1 return f.writeEnvironmentFile(a0, a1, a2, a3) } +type fakeEnvironmentsReader struct { + t testing.TB + getEnvironment func(context.Context, string, string) (sessions.Environment, error) +} + +func (f *fakeEnvironmentsReader) GetEnvironment(a0 context.Context, a1 string, a2 string) (sessions.Environment, error) { + if f.getEnvironment == nil { + unexpectedCall(f.t, "GetEnvironment") + } + return f.getEnvironment(a0, a1, a2) +} + type fakeEnvironments struct { t testing.TB - getEnvironment func(context.Context, string, string) (sessions.Environment, error) authorizeEnvironmentInstallation func(context.Context, identity.Principal, string, string) (string, int64, error) validateEnvironmentInstallation func(context.Context, string, string) (sessions.InstallationAuthorization, error) claimEnvironmentInstallation func(context.Context, string, string, string) error @@ -404,13 +415,6 @@ type fakeEnvironments struct { revokeProjectExecutorCredential func(context.Context, identity.Principal, string, string) error } -func (f *fakeEnvironments) GetEnvironment(a0 context.Context, a1 string, a2 string) (sessions.Environment, error) { - if f.getEnvironment == nil { - unexpectedCall(f.t, "GetEnvironment") - } - return f.getEnvironment(a0, a1, a2) -} - func (f *fakeEnvironments) AuthorizeEnvironmentInstallation(a0 context.Context, a1 identity.Principal, a2 string, a3 string) (string, int64, error) { if f.authorizeEnvironmentInstallation == nil { unexpectedCall(f.t, "AuthorizeEnvironmentInstallation") diff --git a/services/core/internal/api/session_diagnostics_public_compat_test.go b/services/core/internal/api/session_diagnostics_public_compat_test.go index 63694ea55..3f8e5aa7a 100644 --- a/services/core/internal/api/session_diagnostics_public_compat_test.go +++ b/services/core/internal/api/session_diagnostics_public_compat_test.go @@ -67,7 +67,7 @@ func diagnosticRequest(handler http.Handler, path, token string) *httptest.Respo func databaseSessionReads(s *store.Store, pool *pgxpool.Pool) func(*Dependencies, *testFakes) { return func(d *Dependencies, _ *testFakes) { d.Sessions, d.Turns, d.SessionAdmin = s, s, s - d.Items = sessionpg.New(pgunit.NewPool(pool)) + d.Items = sessionpg.New(pgunit.NewPool(pool), nil) } } diff --git a/services/core/internal/execution/archive_cancellation_cleanup_test.go b/services/core/internal/execution/archive_cancellation_cleanup_test.go index 709cefe8e..d3f29bbde 100644 --- a/services/core/internal/execution/archive_cancellation_cleanup_test.go +++ b/services/core/internal/execution/archive_cancellation_cleanup_test.go @@ -113,7 +113,8 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { if scenario.delivery { server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - handler, liveRegistry, err := runtime.NewGateway(s, wsURL) + credentials, heartbeat := testSessions(t, pool, testCredentialCipher(t)) + handler, liveRegistry, err := runtime.NewGateway(credentials, heartbeat, s, wsURL) if err != nil { t.Fatal(err) } @@ -165,7 +166,8 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { t.Fatal("Kill bypassed durable cleanup ownership", allocation, err) } }} - lifecycle := &runtimeLifecycle{store: writer, lease: leased.Lease, registry: registry, config: RuntimeProvider{InstallationID: installation, Provider: provider}, connections: map[string]*runtimeConnection{}} + 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{}} if checkpoint { lifecycle.config.Provider = waitingCleanupCheckpoint{beforeKill: provider.beforeKill} } diff --git a/services/core/internal/execution/device_authority.go b/services/core/internal/execution/device_authority.go index ef1ff6599..fb13e8283 100644 --- a/services/core/internal/execution/device_authority.go +++ b/services/core/internal/execution/device_authority.go @@ -5,17 +5,16 @@ import ( "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" ) // authorizedRuntimePeer fences a connected socket against current credential // authority. A stable device ID does not keep an old credential alive on rotation. -func authorizedRuntimePeer(ctx context.Context, s *store.Store, registry *runtimegateway.Registry, id string) (*runtimegateway.Session, error) { +func authorizedRuntimePeer(ctx context.Context, devices sessions.DeviceReader, registry *runtimegateway.Registry, id string) (*runtimegateway.Session, error) { peer, err := registry.LookupDevice(id) if err != nil { return nil, err } - credential, found, err := s.GetDeviceCredential(ctx, id) + credential, found, err := devices.GetDeviceCredential(ctx, id) if err != nil { return nil, err } @@ -34,5 +33,5 @@ func authorizedRuntimePeer(ctx context.Context, s *store.Store, registry *runtim } func (d *Dispatcher) authorizedPeer(ctx context.Context, id string) (*runtimegateway.Session, error) { - return authorizedRuntimePeer(ctx, d.Store, d.Registry, id) + return authorizedRuntimePeer(ctx, d.SessionsReader, d.Registry, id) } diff --git a/services/core/internal/execution/environment_connections.go b/services/core/internal/execution/environment_connections.go deleted file mode 100644 index ff3be4b66..000000000 --- a/services/core/internal/execution/environment_connections.go +++ /dev/null @@ -1,13 +0,0 @@ -package execution - -import "context" - -// ReplaceEnvironmentConnection begins a serialized transport generation on the owned writer. -func (w *Worker) ReplaceEnvironmentConnection(ctx context.Context, tenant, environment, generation string) error { - return w.dispatcher.Store.ReplaceEnvironmentConnection(ctx, tenant, environment, generation) -} - -// ObserveEnvironmentConnection retains a current transport observation and its event together. -func (w *Worker) ObserveEnvironmentConnection(ctx context.Context, tenant, environment, generation string, revision int64, connected bool) error { - return w.dispatcher.Store.ObserveEnvironmentConnection(ctx, tenant, environment, generation, revision, connected) -} diff --git a/services/core/internal/execution/environment_directory.go b/services/core/internal/execution/environment_directory.go index 8b6c61562..3093a23e0 100644 --- a/services/core/internal/execution/environment_directory.go +++ b/services/core/internal/execution/environment_directory.go @@ -27,7 +27,7 @@ func (r directoryReadRequest) reply(result directoryReadResult) { r.result <- re func (w *Worker) ReadEnvironmentDirectory(ctx context.Context, environment sessions.Environment, path string) (proto.WorkspaceDirectoryResult, error) { ctx, cancel := context.WithTimeout(ctx, 45*time.Second) defer cancel() - current, err := w.admission.GetEnvironment(ctx, environment.TenantID, environment.ID) + current, err := w.dispatcher.SessionsReader.GetEnvironment(ctx, environment.TenantID, environment.ID) if err != nil { return proto.WorkspaceDirectoryResult{}, err } @@ -72,7 +72,7 @@ func (w *Worker) runDirectoryRead(owner context.Context, request directoryReadRe if w.CheckOwnership(check) != nil { return } - environment, err := w.dispatcher.Store.GetEnvironment(check, request.environment.TenantID, request.environment.ID) + environment, err := w.dispatcher.SessionsReader.GetEnvironment(check, request.environment.TenantID, request.environment.ID) if err != nil { result.err = err return @@ -105,7 +105,7 @@ func (w *Worker) runDirectoryRead(owner context.Context, request directoryReadRe } // Capture retains the public Turn after its native Run has been released. prepare := reserved || session.LastTurn != nil && session.LastTurn.ArtifactCaptureStarted - bound, err := w.dispatcher.Store.GetSessionDevice(check, session.TenantID, session.ID) + bound, err := w.dispatcher.SessionsReader.GetSessionDevice(check, session.TenantID, session.ID) if err != nil || !environmentDeviceMatches(session, environment, bound) || !w.directoryDeviceReady(check, bound.ID, session.Engine, placement, prepare) { return } diff --git a/services/core/internal/execution/environment_file_write.go b/services/core/internal/execution/environment_file_write.go index b0b43f194..89e3b379f 100644 --- a/services/core/internal/execution/environment_file_write.go +++ b/services/core/internal/execution/environment_file_write.go @@ -74,7 +74,7 @@ func (w *Worker) runFileWrite(owner context.Context, request fileWriteRequest) f if w.CheckOwnership(ctx) != nil { return unavailable } - environment, err := w.dispatcher.Store.GetEnvironment(ctx, request.environment.TenantID, request.environment.ID) + environment, err := w.dispatcher.SessionsReader.GetEnvironment(ctx, request.environment.TenantID, request.environment.ID) if err != nil { return fileWriteResult{err: err} } @@ -89,7 +89,7 @@ func (w *Worker) runFileWrite(owner context.Context, request fileWriteRequest) f if err != nil { return fileWriteResult{err: err} } - bound, err := w.dispatcher.Store.GetSessionDevice(ctx, session.TenantID, session.ID) + bound, err := w.dispatcher.SessionsReader.GetSessionDevice(ctx, session.TenantID, session.ID) if err != nil || !environmentDeviceMatches(session, environment, bound) || w.dispatcher.Registry == nil { return unavailable } diff --git a/services/core/internal/execution/message_input.go b/services/core/internal/execution/message_input.go index 87e174edc..5a44766d8 100644 --- a/services/core/internal/execution/message_input.go +++ b/services/core/internal/execution/message_input.go @@ -3,6 +3,7 @@ package execution import ( "context" "encoding/json" + 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/sessions" diff --git a/services/core/internal/execution/prepared_dispatch.go b/services/core/internal/execution/prepared_dispatch.go index e8d8d8377..204a3d11b 100644 --- a/services/core/internal/execution/prepared_dispatch.go +++ b/services/core/internal/execution/prepared_dispatch.go @@ -30,7 +30,7 @@ func (d *Dispatcher) RunEnvironmentInput(ctx context.Context, lease Ownership, t if err != nil { return run, err } - environment, err := d.Store.GetSessionEnvironment(ctx, tenantID, sessionID) + environment, err := d.SessionsReader.GetSessionEnvironment(ctx, tenantID, sessionID) if err != nil { return run, err } diff --git a/services/core/internal/execution/runtime_compute.go b/services/core/internal/execution/runtime_compute.go index 858588f7c..835b935ec 100644 --- a/services/core/internal/execution/runtime_compute.go +++ b/services/core/internal/execution/runtime_compute.go @@ -131,7 +131,7 @@ func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.Checkpoint if compute.Status != "running" || !compute.BootstrapComplete { return sandbox.ErrComputeUnconfirmed } - peer, err := authorizedRuntimePeer(ctx, r.store, r.registry, owner.DeviceID) + peer, err := authorizedRuntimePeer(ctx, r.sessions, r.registry, owner.DeviceID) if err != nil { return err } @@ -170,7 +170,7 @@ func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.Checkpoint return err } // Publish disconnected only after receiving the daemon's receipt barrier. - if err := observeRuntimeConnection(ctx, r.store, r.connections, owner.TenantID, owner.EnvironmentID, nil, false); err != nil { + if err := observeRuntimeConnection(ctx, r.sessionExecution, r.connections, owner.TenantID, owner.EnvironmentID, nil, false); err != nil { return err } // The Session-locked phase commit checks pending work and wake requests. diff --git a/services/core/internal/execution/runtime_compute_wake.go b/services/core/internal/execution/runtime_compute_wake.go index a5dbd35f9..339f7847c 100644 --- a/services/core/internal/execution/runtime_compute_wake.go +++ b/services/core/internal/execution/runtime_compute_wake.go @@ -18,7 +18,7 @@ func (r *runtimeLifecycle) wakeCompute(ctx context.Context, p sandbox.Checkpoint return err } } - peer, err := authorizedRuntimePeer(ctx, r.store, r.registry, owner.DeviceID) + peer, err := authorizedRuntimePeer(ctx, r.sessions, r.registry, owner.DeviceID) if err != nil { if !errors.Is(err, sessions.ErrNotFound) && !errors.Is(err, runtimegateway.ErrDeviceNotRegistered) && !errors.Is(err, runtimegateway.ErrSessionClosed) { return err @@ -35,7 +35,7 @@ func (r *runtimeLifecycle) wakeCompute(ctx context.Context, p sandbox.Checkpoint timer := time.NewTicker(100 * time.Millisecond) defer timer.Stop() for { - peer, err = authorizedRuntimePeer(ctx, r.store, r.registry, owner.DeviceID) + peer, err = authorizedRuntimePeer(ctx, r.sessions, r.registry, owner.DeviceID) if err == nil { break } diff --git a/services/core/internal/execution/runtime_connections.go b/services/core/internal/execution/runtime_connections.go index 1b58ed7e9..bce757088 100644 --- a/services/core/internal/execution/runtime_connections.go +++ b/services/core/internal/execution/runtime_connections.go @@ -21,7 +21,7 @@ type runtimeConnection struct { } func (r *runtimeLifecycle) observeConnection(ctx context.Context, owner store.RuntimeAllocation) error { - bound, err := r.store.GetSessionRuntimeDevice(ctx, owner.TenantID, owner.SessionID) + bound, err := r.sessions.GetSessionRuntimeDevice(ctx, owner.TenantID, owner.SessionID) if errors.Is(err, sessions.ErrNotFound) { return nil } @@ -34,21 +34,21 @@ func (r *runtimeLifecycle) observeConnection(ctx context.Context, owner store.Ru if !owner.CreateSettled || owner.State != "running" { return nil } - peer, err := authorizedRuntimePeer(ctx, r.store, r.registry, owner.DeviceID) + peer, err := authorizedRuntimePeer(ctx, r.sessions, r.registry, owner.DeviceID) connected := err == nil if err != nil && !errors.Is(err, sessions.ErrNotFound) && !errors.Is(err, runtimegateway.ErrSessionClosed) && !errors.Is(err, runtimegateway.ErrDeviceNotRegistered) { return err } - return observeRuntimeConnection(ctx, r.store, r.connections, owner.TenantID, owner.EnvironmentID, peer, connected) + return observeRuntimeConnection(ctx, r.sessionExecution, r.connections, owner.TenantID, owner.EnvironmentID, peer, connected) } // Each Environment has one observer: the hosted lifecycle or the Worker loop for // enrolled user compute. Both publish the same durable generation/revision rules. -func observeRuntimeConnection(ctx context.Context, s *store.Store, connections map[string]*runtimeConnection, tenant, environment string, peer *runtimegateway.Session, connected bool) error { +func observeRuntimeConnection(ctx context.Context, operations *sessions.ExecutionOperations, connections map[string]*runtimeConnection, tenant, environment string, peer *runtimegateway.Session, connected bool) error { current := connections[environment] if connected && (current == nil || current.peer != peer) { generation := uuid.NewString() - if err := s.ReplaceEnvironmentConnection(ctx, tenant, environment, generation); err != nil { + if err := operations.ReplaceEnvironmentConnection(ctx, tenant, environment, generation); err != nil { return err } current = &runtimeConnection{peer: peer, generation: generation} @@ -58,7 +58,7 @@ func observeRuntimeConnection(ctx context.Context, s *store.Store, connections m return nil } current.revision++ - if err := s.ObserveEnvironmentConnection(ctx, tenant, environment, current.generation, current.revision, connected); err != nil { + if err := operations.ObserveEnvironmentConnection(ctx, tenant, environment, current.generation, current.revision, connected); err != nil { return err } current.connected = connected @@ -66,7 +66,7 @@ func observeRuntimeConnection(ctx context.Context, s *store.Store, connections m } func (w *Worker) observeEnrolledRuntimes(ctx context.Context) error { - bindings, err := w.dispatcher.Store.ListEnrolledRuntimeBindings(ctx) + bindings, err := w.dispatcher.SessionsReader.ListEnrolledRuntimeBindings(ctx) if err != nil { return err } @@ -78,7 +78,7 @@ func (w *Worker) observeEnrolledRuntimes(ctx context.Context) error { if err != nil && !errors.Is(err, sessions.ErrNotFound) && !errors.Is(err, runtimegateway.ErrSessionClosed) && !errors.Is(err, runtimegateway.ErrDeviceNotRegistered) { return err } - if err := observeRuntimeConnection(ctx, w.dispatcher.Store, w.enrolledConnections, bound.TenantID, bound.EnvironmentID, peer, connected); err != nil { + if err := observeRuntimeConnection(ctx, w.dispatcher.sessionExecution, w.enrolledConnections, bound.TenantID, bound.EnvironmentID, peer, connected); err != nil { if errors.Is(err, sessions.ErrNotFound) || errors.Is(err, sessions.ErrInvalidInput) { continue } diff --git a/services/core/internal/execution/runtime_initialization.go b/services/core/internal/execution/runtime_initialization.go index 5f236e894..4f94d4d52 100644 --- a/services/core/internal/execution/runtime_initialization.go +++ b/services/core/internal/execution/runtime_initialization.go @@ -34,7 +34,7 @@ func (w *Worker) runEnvironmentInitializations(ctx context.Context) error { delete(active, id) case <-ticker.C: } - rows, err := w.dispatcher.Store.ListEnvironmentInitializations(ctx, cursor) + rows, err := w.dispatcher.SessionsReader.ListEnvironmentInitializations(ctx, cursor) if err != nil { return err } @@ -48,7 +48,7 @@ func (w *Worker) runEnvironmentInitializations(ctx context.Context) error { } if owner.State == "running" { // Lost process-local progress cannot prove which side effects ran. - if err := w.dispatcher.Store.FailEnvironmentInitialization(ctx, owner, sessions.ProvisioningFailure{}); err != nil && !errors.Is(err, sessions.ErrNotFound) { + if err := w.dispatcher.sessionExecution.FailEnvironmentInitialization(ctx, owner, sessions.ProvisioningFailure{}); err != nil && !errors.Is(err, sessions.ErrNotFound) { return err } continue @@ -68,12 +68,12 @@ func (w *Worker) runEnvironmentInitializations(ctx context.Context) error { continue } if !found || !harness.Available { - if err := w.dispatcher.Store.FailEnvironmentInitialization(ctx, owner, sessions.ProvisioningFailure{Step: sessions.ProvisioningHarness}); err != nil && !errors.Is(err, sessions.ErrNotFound) { + if err := w.dispatcher.sessionExecution.FailEnvironmentInitialization(ctx, owner, sessions.ProvisioningFailure{Step: sessions.ProvisioningHarness}); err != nil && !errors.Is(err, sessions.ErrNotFound) { return err } continue } - if err := w.dispatcher.Store.ClaimEnvironmentInitialization(ctx, owner); err != nil { + if err := w.dispatcher.sessionExecution.ClaimEnvironmentInitialization(ctx, owner); err != nil { if errors.Is(err, sessions.ErrNotFound) || errors.Is(err, sessions.ErrTurnConflict) { continue } @@ -96,19 +96,19 @@ func (w *Worker) initializeEnvironment(ctx context.Context, owner sessions.Envir failure := sessions.ProvisioningFailure{} err := w.prepareEnvironment(operation, owner, &failure) if err == nil { - err = w.dispatcher.Store.CompleteEnvironmentInitialization(operation, owner) + err = w.dispatcher.sessionExecution.CompleteEnvironmentInitialization(operation, owner) } if err != nil { log.Warn(ctx, "Environment preparation failed", "environment_id", owner.EnvironmentID, "session_id", owner.SessionID) // A later scan settles an unrecorded failure; it never retries the setup. record, stop := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) defer stop() - _ = w.dispatcher.Store.FailEnvironmentInitialization(record, owner, failure) + _ = w.dispatcher.sessionExecution.FailEnvironmentInitialization(record, owner, failure) } } func (w *Worker) prepareEnvironment(ctx context.Context, owner sessions.EnvironmentInitialization, failure *sessions.ProvisioningFailure) error { - environment, err := w.dispatcher.Store.GetEnvironment(ctx, owner.TenantID, owner.EnvironmentID) + environment, err := w.dispatcher.SessionsReader.GetEnvironment(ctx, owner.TenantID, owner.EnvironmentID) if err != nil { return err } @@ -118,7 +118,7 @@ func (w *Worker) prepareEnvironment(ctx context.Context, owner sessions.Environm if json.Unmarshal(environment.Configuration, &cfg) != nil || len(cfg.Files) > 50 { return sessions.ErrInvalidInput } - setup, err := w.dispatcher.Store.ReadEnvironmentSetup(ctx, owner.TenantID, owner.SessionID) + setup, err := w.dispatcher.SessionsReader.ReadEnvironmentSetup(ctx, owner.TenantID, owner.SessionID) if err != nil { return err } @@ -142,7 +142,7 @@ func (w *Worker) prepareEnvironment(ctx context.Context, owner sessions.Environm if err == nil && index < len(cfg.Files) { var metadata environmentconfig.InitialFileMetadata var body []byte - metadata, body, err = w.dispatcher.Store.ReadInitialEnvironmentFile(step, owner.TenantID, owner.SessionID, index) + metadata, body, err = w.dispatcher.SessionsReader.ReadInitialEnvironmentFile(step, owner.TenantID, owner.SessionID, index) if err == nil { err = installInitialFile(step, peer, identity, metadata, body) } diff --git a/services/core/internal/execution/runtime_lifecycle.go b/services/core/internal/execution/runtime_lifecycle.go index dafa48f45..f960edf04 100644 --- a/services/core/internal/execution/runtime_lifecycle.go +++ b/services/core/internal/execution/runtime_lifecycle.go @@ -42,23 +42,25 @@ type RuntimeProvider struct { } type runtimeLifecycle struct { - store *store.Store - 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{} + 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{} } -func newRuntimeManager(owner Owner, deployments *deployment.Service, reader deployment.Reader, registry *runtimegateway.Registry, config *RuntimeProvider) (*runtimeManager, error) { +func newRuntimeManager(owner Owner, deployments *deployment.Service, deploymentReader deployment.Reader, sessionReader sessions.Reader, registry *runtimegateway.Registry, config *RuntimeProvider) (*runtimeManager, error) { if config == nil { return nil, nil } @@ -76,7 +78,7 @@ func newRuntimeManager(owner Owner, deployments *deployment.Service, reader depl } } ctx, stop := context.WithCancel(context.Background()) - return &runtimeManager{store: owner.Store, deployment: owner.Deployment, deploymentService: deployments, deploymentReader: reader, lease: owner.Lease, registry: registry, config: copied, setupInstallationID: config.InstallationID, loadDeployment: config.loadDeployment, prepareDeployment: config.prepareDeployment, publishUnconfigured: config.PublishUnconfigured, setupGate: make(chan struct{}, 1), mutationGate: make(chan struct{}, 1), ctx: ctx, cancel: stop, nodes: make(map[string]*runtimeNode), failed: make(chan error, 1), inventory: make(chan struct{}, 1)}, nil + return &runtimeManager{store: owner.Store, sessions: sessionReader, sessionExecution: owner.Sessions, deployment: owner.Deployment, deploymentService: deployments, deploymentReader: deploymentReader, lease: owner.Lease, registry: registry, config: copied, setupInstallationID: config.InstallationID, loadDeployment: config.loadDeployment, prepareDeployment: config.prepareDeployment, publishUnconfigured: config.PublishUnconfigured, setupGate: make(chan struct{}, 1), mutationGate: make(chan struct{}, 1), ctx: ctx, cancel: stop, nodes: make(map[string]*runtimeNode), failed: make(chan error, 1), inventory: make(chan struct{}, 1)}, nil } func validatedRuntimeProvider(config *RuntimeProvider, registry *runtimegateway.Registry) (RuntimeProvider, error) { @@ -173,7 +175,7 @@ func (r *runtimeLifecycle) provision(ctx context.Context, tenant, environment, p if providerKey != r.config.InstallationID { return store.RuntimeAllocation{}, sandbox.ErrInvalid } - environmentValue, err := r.store.GetEnvironment(ctx, tenant, environment) + environmentValue, err := r.sessions.GetEnvironment(ctx, tenant, environment) if err != nil { return store.RuntimeAllocation{}, err } @@ -396,7 +398,7 @@ func (r *runtimeLifecycle) observe(ctx context.Context, owner store.RuntimeAlloc if err != nil { return err } - environment, err := r.store.GetEnvironment(ctx, owner.TenantID, owner.EnvironmentID) + environment, err := r.sessions.GetEnvironment(ctx, owner.TenantID, owner.EnvironmentID) if err != nil { return err } diff --git a/services/core/internal/execution/runtime_manager.go b/services/core/internal/execution/runtime_manager.go index 95b9acedf..589b1d95b 100644 --- a/services/core/internal/execution/runtime_manager.go +++ b/services/core/internal/execution/runtime_manager.go @@ -9,6 +9,7 @@ import ( "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" ) @@ -18,6 +19,8 @@ var errRuntimeTransition = fmt.Errorf("%w: sandbox configuration is changing", E type runtimeManager struct { store *store.Store + sessions sessions.Reader + sessionExecution *sessions.ExecutionOperations deployment *deployment.ExecutionOperations deploymentService *deployment.Service deploymentReader deployment.Reader @@ -86,7 +89,7 @@ func (m *runtimeManager) node(id string) (*runtimeNode, error) { if n == nil { ctx, stop := context.WithCancel(m.ctx) n = &runtimeNode{lifecycle: &runtimeLifecycle{ - store: m.store, lease: m.lease, registry: m.registry, config: m.config, nodeID: id, + store: m.store, sessions: m.sessions, sessionExecution: m.sessionExecution, lease: m.lease, registry: m.registry, config: m.config, nodeID: id, gate: make(chan struct{}, 1), ctx: ctx, stop: stop, connections: make(map[string]*runtimeConnection), wakeHints: make(chan struct{}, 1), }} diff --git a/services/core/internal/execution/runtime_wake_hint.go b/services/core/internal/execution/runtime_wake_hint.go index 852cf7be1..074ce4c59 100644 --- a/services/core/internal/execution/runtime_wake_hint.go +++ b/services/core/internal/execution/runtime_wake_hint.go @@ -23,7 +23,7 @@ func (w *Worker) hintRuntimeWake(ctx context.Context, session sessions.Session) } lookup, cancel := context.WithTimeout(ctx, time.Second) defer cancel() - environment, err := w.admission.GetSessionEnvironment(lookup, session.TenantID, session.ID) + environment, err := w.dispatcher.SessionsReader.GetSessionEnvironment(lookup, session.TenantID, session.ID) if err != nil || environment.Initialization != "complete" { return } diff --git a/services/core/internal/execution/sandbox_deployment_drain_test.go b/services/core/internal/execution/sandbox_deployment_drain_test.go index 5a66a587a..ddab4df33 100644 --- a/services/core/internal/execution/sandbox_deployment_drain_test.go +++ b/services/core/internal/execution/sandbox_deployment_drain_test.go @@ -94,7 +94,7 @@ func testLifecycleCancellationPreservesLease(t *testing.T, mode string) { defer hub.Close() id := uuid.NewString() configuration := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)} - m, err := newRuntimeManager(owner, deployments, reader, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return configuration, nil })) + m, err := newRuntimeManager(owner, deployments, reader, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return configuration, nil })) if err != nil { t.Fatal(err) } @@ -232,7 +232,7 @@ func TestSandboxDeploymentDrainFailureCannotReactivate(t *testing.T) { defer hub.Close() id := uuid.NewString() configuration := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)} - m, err := newRuntimeManager(owner, deployments, reader, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return configuration, nil })) + m, err := newRuntimeManager(owner, deployments, reader, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return configuration, nil })) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_deployment_setup_test.go b/services/core/internal/execution/sandbox_deployment_setup_test.go index 99045edab..35271c554 100644 --- a/services/core/internal/execution/sandbox_deployment_setup_test.go +++ b/services/core/internal/execution/sandbox_deployment_setup_test.go @@ -25,7 +25,7 @@ func TestDeferredSandboxDeploymentLoadsOnceBeforeNodeCreation(t *testing.T) { var selected atomic.Bool var loads atomic.Int32 configuration := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)} - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { loads.Add(1) if !selected.Load() { return nil, nil @@ -70,7 +70,7 @@ func TestDeferredSandboxDeploymentLoadsOnceBeforeNodeCreation(t *testing.T) { func TestDeferredSandboxDeploymentShutdownCancelsLoad(t *testing.T) { entered := make(chan struct{}) - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(uuid.NewString(), func(ctx context.Context) (*RuntimeProvider, error) { + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(uuid.NewString(), func(ctx context.Context) (*RuntimeProvider, error) { close(entered) <-ctx.Done() return nil, ctx.Err() @@ -95,7 +95,7 @@ func TestDeferredSandboxProviderFailureKeepsRecoveryAvailable(t *testing.T) { available := false loadErr := ErrExecutionUnavailable configuration := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)} - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { if !available { return nil, loadErr } @@ -127,7 +127,7 @@ func TestRejectedSandboxCandidatePreservesActiveGeneration(t *testing.T) { id := uuid.NewString() config := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)} rejected := errors.New("candidate provider unavailable") - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, unitDeploymentService(t), nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil }, + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, unitDeploymentService(t), nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil }, func(context.Context, deployment.Setup) (PreparedRuntimeDeployment, error) { return PreparedRuntimeDeployment{}, rejected })) @@ -159,7 +159,7 @@ func TestRejectedSandboxCandidatePreservesActiveGeneration(t *testing.T) { func TestSandboxCandidateValidationDoesNotHoldManagerLock(t *testing.T) { id := uuid.NewString() entered, release := make(chan struct{}), make(chan struct{}) - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, unitDeploymentService(t), nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil }, + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, unitDeploymentService(t), nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil }, func(context.Context, deployment.Setup) (PreparedRuntimeDeployment, error) { close(entered) <-release @@ -192,7 +192,7 @@ func TestCommittedSandboxCandidatePublishesAfterShutdown(t *testing.T) { hub := node.NewHub(node.HubOptions{}) defer hub.Close() id := uuid.NewString() - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_deployment_switch_test.go b/services/core/internal/execution/sandbox_deployment_switch_test.go index 38d6375d9..2d31d3d32 100644 --- a/services/core/internal/execution/sandbox_deployment_switch_test.go +++ b/services/core/internal/execution/sandbox_deployment_switch_test.go @@ -19,7 +19,7 @@ func TestSandboxManagerSwitchDrainsBeforeDirectActivation(t *testing.T) { defer hub.Close() id := uuid.NewString() config := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)} - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) if err != nil { t.Fatal(err) } @@ -80,7 +80,7 @@ func TestSandboxManagerSwitchDrainsBeforeDirectActivation(t *testing.T) { func TestSandboxManagerFailedActivationStaysPaused(t *testing.T) { id := uuid.NewString() - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, errors.New("provider unavailable") })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, errors.New("provider unavailable") })) if err != nil { t.Fatal(err) } @@ -101,7 +101,7 @@ func TestSandboxManagerCancelledSwitchCannotResumeBeforeDrain(t *testing.T) { defer hub.Close() id := uuid.NewString() config := &RuntimeProvider{InstallationID: id, ProviderKind: "docker", Mode: "nodes", Generation: 1, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)} - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) if err != nil { t.Fatal(err) } @@ -155,7 +155,7 @@ func TestSandboxActivationCannotBypassOutstandingDrain(t *testing.T) { hub := node.NewHub(node.HubOptions{}) defer hub.Close() id := uuid.NewString() - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil }, func(_ context.Context, setup deployment.Setup) (PreparedRuntimeDeployment, error) { return PreparedRuntimeDeployment{Config: &RuntimeProvider{InstallationID: setup.InstallationID, ProviderKind: setup.Provider, Mode: setup.Mode, CoreURL: "https://core.example/api/v1", BackendFingerprint: setup.BackendFingerprint, Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)}}, nil diff --git a/services/core/internal/execution/sandbox_generations_test.go b/services/core/internal/execution/sandbox_generations_test.go index 5a179a8d6..4ae75b9e4 100644 --- a/services/core/internal/execution/sandbox_generations_test.go +++ b/services/core/internal/execution/sandbox_generations_test.go @@ -50,7 +50,7 @@ func TestE2BReplacementVerifiesTwiceAndNeverPublishesFailedCommit(t *testing.T) }, FenceCredential: func(context.Context) (func(), error) { fenced++; return func() { released++ }, nil }, Publish: func(*RuntimeProvider) { published++ }}, nil }) - m, err := newRuntimeManager(owner, deployments, reader, runtimegateway.NewRegistry(), config) + m, err := newRuntimeManager(owner, deployments, reader, nil, runtimegateway.NewRegistry(), config) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_provider_contract_test.go b/services/core/internal/execution/sandbox_provider_contract_test.go index 0c154fe74..531b3fbbc 100644 --- a/services/core/internal/execution/sandbox_provider_contract_test.go +++ b/services/core/internal/execution/sandbox_provider_contract_test.go @@ -18,7 +18,7 @@ func TestSandboxProviderRegistrationDoesNotRequireAnExecutionVendorBranch(t *tes id := uuid.NewString() config := &RuntimeProvider{InstallationID: id, ProviderKind: "contract-fixture", Mode: mode, CoreURL: "https://core.example/api/v1", BackendFingerprint: strings.Repeat("a", 64), Provider: &lifecycleOnlySandbox{}} - m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) + m, err := newRuntimeManager(Owner{Lease: heldLease{}}, nil, nil, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return config, nil })) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_reset_test.go b/services/core/internal/execution/sandbox_reset_test.go index 7bcc22cfa..da24eec29 100644 --- a/services/core/internal/execution/sandbox_reset_test.go +++ b/services/core/internal/execution/sandbox_reset_test.go @@ -46,17 +46,25 @@ func resetManagerStoreConfig(t *testing.T, configure func(*pgxpool.Config)) (*st func resetManagerStoreDB(t *testing.T, configure func(*pgxpool.Config)) (*store.Store, Owner, *deployment.Service, deployment.Reader, *pgxpool.Pool) { t.Helper() pool := pgtest.OpenIsolated(t, configure) + cipher := testCredentialCipher(t) + s := store.NewWithCredentialCipher(pool, cipher) + owner, deployments, reader := testOwner(t, pool, cipher, s) + return s, owner, deployments, reader, pool +} + +// testCredentialCipher is the credential key of the Store resetManagerStoreDB +// builds, for adapters built beside it. +func testCredentialCipher(t *testing.T) *credentialcrypto.Cipher { + t.Helper() cipher, err := credentialcrypto.New(bytes.Repeat([]byte{8}, 32)) if err != nil { t.Fatal(err) } - s := store.NewWithCredentialCipher(pool, cipher) - owner, deployments, reader := testOwner(t, pool, cipher, s) - return s, owner, deployments, reader, pool + return cipher } // testOwner acquires the execution lease on pool and builds s's execution -// writer and the deployment execution operations on it, and the pooled +// writer and the deployment and Session execution operations on it, and the pooled // deployment service and reader, as cmd/server does. The lease closes when the // test ends. func testOwner(t *testing.T, pool *pgxpool.Pool, cipher *credentialcrypto.Cipher, s *store.Store) (Owner, *deployment.Service, deployment.Reader) { @@ -98,7 +106,7 @@ func TestSandboxResetPageTimeoutRecoversCommittedOwner(t *testing.T) { } return &RuntimeProvider{InstallationID: id, ProviderKind: setup.Provider, Mode: setup.Mode, Generation: setup.Generation, CoreURL: "https://core.example/api/v1", BackendFingerprint: setup.BackendFingerprint, Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)}, nil }) - m, err := newRuntimeManager(owner, deployments, reader, runtimegateway.NewRegistry(), config) + m, err := newRuntimeManager(owner, deployments, reader, nil, runtimegateway.NewRegistry(), config) if err != nil { t.Fatal(err) } @@ -219,7 +227,7 @@ func TestSandboxResetPublishesCommittedGenerationWithoutReading(t *testing.T) { } published = append(published, generation) } - m, err := newRuntimeManager(owner, deployments, adapter, runtimegateway.NewRegistry(), config) + m, err := newRuntimeManager(owner, deployments, adapter, nil, runtimegateway.NewRegistry(), config) if err != nil { t.Fatal(err) } @@ -254,7 +262,7 @@ func TestCommittedResetViewStopsOwnerWithoutLease(t *testing.T) { return deployment.Snapshot{Record: deployment.Record{InstallationID: id, WebManaged: true, Generation: 1}}, nil }} deployments, operations := deploymentOperations(t, &strictDeploymentStorage{t: t}, reader, &strictExecutionStorage{t: t}) - m, err := newRuntimeManager(Owner{Lease: lostLease{}, Deployment: operations}, deployments, reader, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })) + m, err := newRuntimeManager(Owner{Lease: lostLease{}, Deployment: operations}, deployments, reader, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })) if err != nil { t.Fatal(err) } @@ -276,7 +284,7 @@ func TestCommittedResetViewStopsOwnerWithoutLease(t *testing.T) { func TestSandboxResetChangesReturnViewReadAfterCommit(t *testing.T) { _, owner, deployments, reader := resetManagerStore(t) id := initializeE2BDeployment(t, owner) - m, err := newRuntimeManager(owner, deployments, reader, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })) + m, err := newRuntimeManager(owner, deployments, reader, nil, runtimegateway.NewRegistry(), NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_snapshot_budget_test.go b/services/core/internal/execution/sandbox_snapshot_budget_test.go index 9411356cb..152df9d6d 100644 --- a/services/core/internal/execution/sandbox_snapshot_budget_test.go +++ b/services/core/internal/execution/sandbox_snapshot_budget_test.go @@ -74,7 +74,7 @@ func TestSandboxResetSnapshotFitsPageBudget(t *testing.T) { } return &RuntimeProvider{InstallationID: id, ProviderKind: setup.Provider, Generation: setup.Generation, Mode: setup.Mode, CoreURL: "https://core.example/api/v1", BackendFingerprint: setup.BackendFingerprint, Provider: hub.Proxy(uuid.NewString(), "docker", docker.Operations(), 1)}, nil }) - m, err := newRuntimeManager(owner, deployments, reader, runtimegateway.NewRegistry(), configuration) + m, err := newRuntimeManager(owner, deployments, reader, nil, runtimegateway.NewRegistry(), configuration) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/session_fixture_test.go b/services/core/internal/execution/session_fixture_test.go new file mode 100644 index 000000000..8e6f3e322 --- /dev/null +++ b/services/core/internal/execution/session_fixture_test.go @@ -0,0 +1,24 @@ +package execution + +import ( + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" + "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" +) + +// testSessions builds the pooled Session adapter on pool with cipher and the +// Session service on it, as cmd/server does. +func testSessions(t *testing.T, pool *pgxpool.Pool, cipher *credentialcrypto.Cipher) (*sessionpg.Store, *sessions.Service) { + t.Helper() + adapter := sessionpg.New(pgunit.NewPool(pool), cipher) + service, err := sessions.NewService(adapter) + if err != nil { + t.Fatal(err) + } + return adapter, service +} diff --git a/services/core/internal/execution/worker.go b/services/core/internal/execution/worker.go index f42cd8589..2f584c543 100644 --- a/services/core/internal/execution/worker.go +++ b/services/core/internal/execution/worker.go @@ -77,7 +77,7 @@ func StartWorker(ctx context.Context, dispatcher *Dispatcher, owner Owner) (_ *W } owned.notifications = &executionNotifications{} worker := &Worker{concurrency: dispatcher.MaxConcurrentExecutions, dispatcher: owned, admission: dispatcher.Store, lease: owner.Lease, directoryReads: make(chan directoryReadRequest), fileWrites: make(chan fileWriteRequest), stopped: make(chan struct{}), scheduleWake: make(chan struct{}, 1), enrolledConnections: make(map[string]*runtimeConnection)} - worker.runtimes, err = newRuntimeManager(owner, owned.Deployment, owned.DeploymentReader, owned.Registry, owned.ManagedRuntimes) + worker.runtimes, err = newRuntimeManager(owner, owned.Deployment, owned.DeploymentReader, owned.SessionsReader, owned.Registry, owned.ManagedRuntimes) if err != nil { return nil, err } @@ -104,7 +104,7 @@ func StartWorker(ctx context.Context, dispatcher *Dispatcher, owner Owner) (_ *W if err != nil { return nil, err } - if err = owned.Store.ReconcileEnvironmentConnections(ctx); err != nil { + if err = owned.sessionExecution.ReconcileEnvironmentConnections(ctx); err != nil { return nil, err } if err = worker.reconcile(ctx); err != nil { diff --git a/services/core/internal/execution/worker_device.go b/services/core/internal/execution/worker_device.go index 6999868e9..b17d83eb7 100644 --- a/services/core/internal/execution/worker_device.go +++ b/services/core/internal/execution/worker_device.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "errors" + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -73,7 +74,7 @@ func (w *Worker) bindSessionDevice(ctx context.Context, session sessions.Session return false, sessions.ErrInvalidInput } if snapshot.Environment != nil && (snapshot.Environment.Type == "openai_hosted" || snapshot.Environment.Type == "self_hosted") { - environment, err := w.dispatcher.Store.GetSessionEnvironment(ctx, session.TenantID, session.ID) + environment, err := w.dispatcher.SessionsReader.GetSessionEnvironment(ctx, session.TenantID, session.ID) if err != nil { return false, err } @@ -92,13 +93,13 @@ func (w *Worker) bindSessionDevice(ctx context.Context, session sessions.Session // resumes. Its first control frame must remain the lifecycle's Resume. return false, nil } - bound, err := w.dispatcher.Store.GetSessionDevice(ctx, session.TenantID, session.ID) + bound, err := w.dispatcher.SessionsReader.GetSessionDevice(ctx, session.TenantID, session.ID) if errors.Is(err, sessions.ErrNotFound) { return false, nil } return err == nil && environmentDeviceMatches(session, environment, bound) && ready(bound.ID), err } - bound, err := w.dispatcher.Store.GetSessionDevice(ctx, session.TenantID, session.ID) + bound, err := w.dispatcher.SessionsReader.GetSessionDevice(ctx, session.TenantID, session.ID) if err == nil { return bound.EnvironmentID == "" && ready(bound.ID), nil } @@ -113,7 +114,7 @@ func (w *Worker) bindSessionDevice(ctx context.Context, session sessions.Session if !ready(device.ID) { continue } - err := w.dispatcher.Store.BindSessionDevice(ctx, session.TenantID, session.ID, device.ID) + err := w.dispatcher.sessionExecution.BindSessionDevice(ctx, session.TenantID, session.ID, device.ID) return err == nil, err } return false, nil diff --git a/services/core/internal/persistence/postgres/sessionpg/artifacts_test.go b/services/core/internal/persistence/postgres/sessionpg/artifacts_test.go index abdded5e0..0c06e8f94 100644 --- a/services/core/internal/persistence/postgres/sessionpg/artifacts_test.go +++ b/services/core/internal/persistence/postgres/sessionpg/artifacts_test.go @@ -67,7 +67,7 @@ func stagedRows(t *testing.T, pool *pgxpool.Pool, turn string) int { func stagingService(t *testing.T, pool *pgxpool.Pool) (*Store, *sessions.Service) { t.Helper() - store := New(pgunit.NewPool(pool)) + store := New(pgunit.NewPool(pool), nil) service, err := sessions.NewService(store) if err != nil { t.Fatal(err) diff --git a/services/core/internal/persistence/postgres/sessionpg/devices.go b/services/core/internal/persistence/postgres/sessionpg/devices.go new file mode 100644 index 000000000..30b04e7f1 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/devices.go @@ -0,0 +1,225 @@ +package sessionpg + +import ( + "context" + "errors" + "fmt" + + "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/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +func (s *Store) GetSessionDevice(ctx context.Context, tenant, session string) (sessions.ExecutionDevice, error) { + lookup, err := DeviceLookup(tenant, session) + if err != nil { + return sessions.ExecutionDevice{}, err + } + ready, err := s.units.Queries().GetSessionInitializationReady(ctx, sqlc.GetSessionInitializationReadyParams(lookup)) + if errors.Is(err, pgx.ErrNoRows) || (err == nil && !ready) { + return sessions.ExecutionDevice{}, sessions.ErrNotFound + } + if err != nil { + return sessions.ExecutionDevice{}, err + } + return s.GetSessionRuntimeDevice(ctx, tenant, session) +} + +func (s *Store) GetSessionRuntimeDevice(ctx context.Context, tenant, session string) (sessions.ExecutionDevice, error) { + lookup, err := DeviceLookup(tenant, session) + if err != nil { + return sessions.ExecutionDevice{}, err + } + device, found, err := loadSessionDevice(ctx, s.units.Queries(), lookup.TenantID, lookup.ID) + if err == nil && !found { + return sessions.ExecutionDevice{}, sessions.ErrNotFound + } + return device, err +} + +// loadSessionDevice reads on q the device bound to the tenant's Session and +// reports whether the Session has one that was not revoked. +func loadSessionDevice(ctx context.Context, q *sqlc.Queries, tenant, session pgtype.UUID) (sessions.ExecutionDevice, bool, error) { + row, err := q.GetSessionDevice(ctx, sqlc.GetSessionDeviceParams{TenantID: tenant, ID: session}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.ExecutionDevice{}, false, nil + } + if err != nil { + return sessions.ExecutionDevice{}, false, err + } + return sessions.ExecutionDevice{ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name, EnvironmentID: optionalID(row.EnvironmentID)}, true, nil +} + +// GetDeviceCredential reads a device's credential for the Runtime gateway. A +// malformed or unknown device has no credential. The standalone service +// assigns no product WorkspaceID. +func (s *Store) GetDeviceCredential(ctx context.Context, device string) (runtimedevice.Credential, bool, error) { + id, err := pgunit.ParseID(device) + if err != nil { + return runtimedevice.Credential{}, false, nil + } + row, err := s.units.Queries().GetDeviceCredential(ctx, id) + if errors.Is(err, pgx.ErrNoRows) { + return runtimedevice.Credential{}, false, nil + } + if err != nil { + return runtimedevice.Credential{}, false, err + } + return runtimedevice.Credential{ + ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name, Type: runtimedevice.RuntimeTypeAgentDaemon, + CredentialHash: row.CredentialHash, RuntimeNodeID: row.RuntimeNodeID, RuntimeAllocationID: row.RuntimeAllocationID, + }, true, nil +} + +func (s *Store) ListEnrolledRuntimeBindings(ctx context.Context) ([]sessions.EnrolledRuntimeBinding, error) { + rows, err := s.units.Queries().ListEnrolledRuntimeBindings(ctx) + if err != nil { + return nil, err + } + result := make([]sessions.EnrolledRuntimeBinding, 0, len(rows)) + for _, row := range rows { + result = append(result, sessions.EnrolledRuntimeBinding{ + DeviceID: optionalID(row.DeviceID), TenantID: optionalID(row.TenantID), + EnvironmentID: optionalID(row.EnvironmentID), SessionID: optionalID(row.SessionID), + }) + } + return result, nil +} + +func (s *Store) CreateDevice(ctx context.Context, tenant string, registration sessions.DeviceRegistration) (sessions.ExecutionDevice, error) { + tenantID, err := parseID(tenant) + if err != nil { + return sessions.ExecutionDevice{}, err + } + id, err := s.units.Queries().CreateDevice(ctx, sqlc.CreateDeviceParams{ + ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, TenantID: tenantID, + Name: registration.Name, CredentialHash: pgtype.Text{String: registration.CredentialHash, Valid: true}, + }) + if err != nil { + return sessions.ExecutionDevice{}, fmt.Errorf("create execution device: %w", err) + } + return sessions.ExecutionDevice{ID: uuid.UUID(id.Bytes).String(), Name: registration.Name}, nil +} + +func (s *Store) RevokeDevice(ctx context.Context, tenant, device string) error { + lookup, err := DeviceLookup(tenant, device) + if err != nil { + return err + } + n, err := s.units.Queries().RevokeDevice(ctx, sqlc.RevokeDeviceParams(lookup)) + if err == nil && n == 0 { + return sessions.ErrNotFound + } + return err +} + +func (s *Store) TouchDevice(ctx context.Context, device string) (bool, error) { + id, err := parseID(device) + if err != nil { + return false, err + } + n, err := s.units.Queries().TouchDevice(ctx, id) + return n > 0, err +} + +func (s *Store) TouchAuthenticatedDevice(ctx context.Context, device, credentialHash string) (bool, error) { + id, err := parseID(device) + if err != nil { + return false, err + } + n, err := s.units.Queries().TouchAuthenticatedDevice(ctx, sqlc.TouchAuthenticatedDeviceParams{ID: id, CredentialHash: credentialHash}) + return n > 0, err +} + +// WithEnrollment authenticates the credential and reads the Environment on +// the pool before it knows which Session to lock; the enrollment transaction +// rechecks that authority under the Session lock. +func (s *Store) WithEnrollment(ctx context.Context, environment, credentialHash string, apply func(context.Context, sessions.EnrollmentTx, sessions.Environment, sessions.LockedSession) error) error { + tenant, err := s.AuthenticateEnvironmentExecutor(ctx, environment, credentialHash) + if err != nil { + return err + } + current, err := s.GetEnvironment(ctx, tenant, environment) + if err != nil { + return err + } + lookup, err := DeviceLookup(tenant, current.ID) + if err != nil { + return err + } + return WithSession(ctx, s.units, lookup.TenantID, pgunit.PathID(current.SessionID), func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + tx := &enrollmentTx{SessionTx: BindSession(q, lookup.TenantID, pgunit.PathID(current.SessionID)), environment: lookup.ID, credentialHash: credentialHash} + return apply(ctx, tx, current, locked) + }) +} + +// enrollmentTx is one enrollment inside the transaction of the Environment's +// Session. +type enrollmentTx struct { + *SessionTx + environment pgtype.UUID + credentialHash string +} + +var _ sessions.EnrollmentTx = (*enrollmentTx)(nil) + +func (t *enrollmentTx) AuthorizeEnrollment(ctx context.Context) (sessions.EnrollmentAuthority, error) { + row, err := t.q.AuthorizeRuntimeEnrollment(ctx, sqlc.AuthorizeRuntimeEnrollmentParams{EnvironmentID: t.environment, TenantID: t.tenant, TokenSha256: t.credentialHash}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.EnrollmentAuthority{}, sessions.ErrNotFound + } + if err != nil { + return sessions.EnrollmentAuthority{}, err + } + return sessions.EnrollmentAuthority{KeyID: optionalID(row.KeyID), WorkspaceDirectory: row.WorkspaceDirectory}, nil +} + +func (t *enrollmentTx) EnrollDevice(ctx context.Context, key string) (string, error) { + keyID, err := parseID(key) + if err != nil { + return "", err + } + row, err := t.q.EnrollRuntimeDevice(ctx, sqlc.EnrollRuntimeDeviceParams{ + ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, TenantID: t.tenant, EnvironmentID: t.environment, ExecutorKeyID: keyID, + }) + if errors.Is(err, pgx.ErrNoRows) { + return "", sessions.ErrDeviceBindingConflict + } + if err != nil { + return "", err + } + return uuid.UUID(row.ID.Bytes).String(), nil +} + +var _ sessions.DeviceBindingTx = (*SessionTx)(nil) + +func (t *SessionTx) LoadDevice(ctx context.Context, device string) (bool, error) { + id, err := parseID(device) + if err != nil { + return false, err + } + _, err = t.q.GetDevice(ctx, sqlc.GetDeviceParams{TenantID: t.tenant, ID: id}) + if errors.Is(err, pgx.ErrNoRows) { + return false, nil + } + return err == nil, err +} + +// BindDevice binds the device to the Session; a Session bound to another +// device is sessions.ErrDeviceBindingConflict. +func (t *SessionTx) BindDevice(ctx context.Context, device string) error { + id, err := parseID(device) + if err != nil { + return err + } + _, err = t.q.BindSessionDevice(ctx, sqlc.BindSessionDeviceParams{TenantID: t.tenant, ID: t.session, ID_2: id}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.ErrDeviceBindingConflict + } + return err +} diff --git a/services/core/internal/persistence/postgres/sessionpg/devices_test.go b/services/core/internal/persistence/postgres/sessionpg/devices_test.go new file mode 100644 index 000000000..71c17500a --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/devices_test.go @@ -0,0 +1,71 @@ +package sessionpg + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// newDevice stores a device of tenant, dedicated to environment when it is +// valid, and returns its ID. +func newDevice(t *testing.T, pool *pgxpool.Pool, tenant, environment pgtype.UUID) string { + t.Helper() + id := uuid.NewString() + exec(t, pool, `INSERT INTO devices(id, tenant_id, name, credential_hash, environment_id) VALUES ($1, $2, 'runtime', $3, $4)`, id, tenant, strings.Repeat("a", 64), environment) + return id +} + +func TestBindSessionDeviceTranslatesTheBindingOutcome(t *testing.T) { + pool := pgtest.OpenIsolated(t, nil) + lease, err := pgunit.AcquireLease(t.Context(), pool) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = lease.Close(context.Background()) }) + operations, err := sessions.NewExecutionOperations(NewExecution(lease)) + if err != nil { + t.Fatal(err) + } + tenantID, sessionID, _ := newEnvironment(t, pool, "self_hosted", "pending") + tenant, session := uuidText(tenantID), uuidText(sessionID) + _, _, otherEnvironment := newEnvironment(t, pool, "self_hosted", "pending") + otherTenant, _, _ := newEnvironment(t, pool, "self_hosted", "pending") + device, replacement := newDevice(t, pool, tenantID, pgtype.UUID{}), newDevice(t, pool, tenantID, pgtype.UUID{}) + revoked := newDevice(t, pool, tenantID, pgtype.UUID{}) + exec(t, pool, `UPDATE devices SET revoked_at = clock_timestamp() WHERE id = $1`, revoked) + + for name, test := range map[string]struct { + device string + want error + }{ + "another tenant's device": {newDevice(t, pool, otherTenant, pgtype.UUID{}), sessions.ErrNotFound}, + "revoked device": {revoked, sessions.ErrNotFound}, + "device of another Environment": {newDevice(t, pool, tenantID, otherEnvironment), sessions.ErrDeviceBindingConflict}, + "malformed device": {"device", sessions.ErrInvalidInput}, + } { + if err := operations.BindSessionDevice(t.Context(), tenant, session, test.device); !errors.Is(err, test.want) { + t.Fatalf("%s: %v, want %v", name, err, test.want) + } + } + for range 2 { + if err := operations.BindSessionDevice(t.Context(), tenant, session, device); err != nil { + t.Fatal("binding the bound device again", err) + } + } + if err := operations.BindSessionDevice(t.Context(), tenant, session, replacement); !errors.Is(err, sessions.ErrDeviceBindingConflict) { + t.Fatal("rebinding the Session to another device", err) + } + var bound string + if err := pool.QueryRow(t.Context(), `SELECT device_id::text FROM session_devices WHERE session_id = $1`, sessionID).Scan(&bound); err != nil || bound != device { + t.Fatal("Session bound to", bound, err) + } +} diff --git a/services/core/internal/persistence/postgres/sessionpg/environment.go b/services/core/internal/persistence/postgres/sessionpg/environment.go index fd6667bfe..bf290d46c 100644 --- a/services/core/internal/persistence/postgres/sessionpg/environment.go +++ b/services/core/internal/persistence/postgres/sessionpg/environment.go @@ -9,9 +9,92 @@ 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/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) +// LoadEnvironment reads, on q, the tenant's Environment of a Session that was +// not publicly deleted. A malformed tenant is sessions.ErrInvalidInput, and a +// malformed Environment ID a missing Environment, sessions.ErrNotFound. Other +// adapters read the Environment of their operation with it. +func LoadEnvironment(ctx context.Context, q *sqlc.Queries, tenant, environment string) (sessions.Environment, error) { + tenantID, err := parseID(tenant) + if err != nil { + return sessions.Environment{}, err + } + row, err := q.GetEnvironment(ctx, sqlc.GetEnvironmentParams{TenantID: tenantID, ID: pgunit.PathID(environment)}) + return EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, err) +} + +// LoadSessionEnvironment reads, on q, the Environment of the tenant's Session +// that was not publicly deleted, or sessions.ErrNotFound. A malformed ID is +// sessions.ErrInvalidInput. Other adapters read the Environment of their +// operation with it. +func LoadSessionEnvironment(ctx context.Context, q *sqlc.Queries, tenant, session string) (sessions.Environment, error) { + tenantID, err := parseID(tenant) + if err != nil { + return sessions.Environment{}, err + } + id, err := parseID(session) + if err != nil { + return sessions.Environment{}, err + } + row, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: tenantID, ID: id}) + return EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, err) +} + +// CreateSessionEnvironment creates, on q, the Environment of a new Session +// whose configuration has one, as sessions.CreatesEnvironment decides. Its +// preparation is pending when the Session froze a setup or initial files, and +// complete otherwise, so it runs after those are saved. +func CreateSessionEnvironment(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { + return q.CreateEnvironment(ctx, sqlc.CreateEnvironmentParams{ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, SessionID: session}) +} + +func (s *Store) GetEnvironment(ctx context.Context, tenant, environment string) (sessions.Environment, error) { + return LoadEnvironment(ctx, s.units.Queries(), tenant, environment) +} + +func (s *Store) GetSessionEnvironment(ctx context.Context, tenant, session string) (sessions.Environment, error) { + return LoadSessionEnvironment(ctx, s.units.Queries(), tenant, session) +} + +func (s *Store) ListEnvironmentInitializations(ctx context.Context, after string) ([]sessions.EnvironmentInitialization, error) { + id, err := pageAfter(after) + if err != nil { + return nil, err + } + rows, err := s.units.Queries().ListEnvironmentInitializations(ctx, id) + if err != nil { + return nil, err + } + result := make([]sessions.EnvironmentInitialization, 0, len(rows)) + for _, row := range rows { + result = append(result, sessions.EnvironmentInitialization{ + EnvironmentID: optionalID(row.ID), SessionID: optionalID(row.SessionID), TenantID: optionalID(row.TenantID), + DeviceID: optionalID(row.DeviceID), State: row.Initialization, Engine: row.Engine, + }) + } + return result, nil +} + +// pageAfter parses the Environment a page of Environments starts after; empty +// starts from the first one. +func pageAfter(after string) (pgtype.UUID, error) { + if after == "" { + return pgtype.UUID{Valid: true}, nil + } + return parseID(after) +} + +// optionalID formats a nullable identifier, empty when it is NULL. +func optionalID(id pgtype.UUID) string { + if !id.Valid { + return "" + } + return uuid.UUID(id.Bytes).String() +} + // LoadEnvironmentInput reads the Session's latest Environment input // reservation that no Turn has admitted or superseded, with the facts of its // Environment, and nil when there is none. sessions.InputActivity projects it. diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_environment.go b/services/core/internal/persistence/postgres/sessionpg/execution_environment.go new file mode 100644 index 000000000..505f859ec --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution_environment.go @@ -0,0 +1,164 @@ +package sessionpg + +import ( + "context" + "errors" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + + "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/sessions" +) + +func (e *Execution) WithInitialization(ctx context.Context, tenant, session string, apply func(context.Context, sessions.InitializationTx, sessions.LockedSession) error) error { + tenantID, err := parseID(tenant) + if err != nil { + return err + } + sessionID := pgunit.PathID(session) + return WithSession(ctx, e.lease, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + return apply(ctx, &environmentTx{SessionTx: BindSession(q, tenantID, sessionID)}, locked) + }) +} + +// WithConnection reads the Environment in the transaction before it locks +// the Environment's Session. +func (e *Execution) WithConnection(ctx context.Context, tenant, environment string, apply func(context.Context, sessions.ConnectionTx, sessions.LockedSession) error) error { + return e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + current, err := LoadEnvironment(ctx, q, tenant, environment) + if err != nil { + return err + } + lookup, err := DeviceLookup(tenant, current.SessionID) + if err != nil { + return err + } + return inSession(ctx, q, lookup.TenantID, lookup.ID, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + return apply(ctx, &environmentTx{SessionTx: BindSession(q, lookup.TenantID, lookup.ID)}, locked) + }) + }) +} + +func (e *Execution) WithDeviceBinding(ctx context.Context, tenant, session string, apply func(context.Context, sessions.DeviceBindingTx, sessions.LockedSession) error) error { + lookup, err := DeviceLookup(tenant, session) + if err != nil { + return err + } + return WithSession(ctx, e.lease, lookup.TenantID, lookup.ID, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + return apply(ctx, BindSession(q, lookup.TenantID, lookup.ID), locked) + }) +} + +func (e *Execution) ListEnvironmentConnections(ctx context.Context, after string) ([]sessions.EnvironmentKey, error) { + id, err := pageAfter(after) + if err != nil { + return nil, err + } + var rows []sqlc.ListEnvironmentConnectionsRow + err = e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + var err error + rows, err = sqlc.New(tx).ListEnvironmentConnections(ctx, id) + return err + }) + if err != nil { + return nil, err + } + keys := make([]sessions.EnvironmentKey, 0, len(rows)) + for _, row := range rows { + keys = append(keys, sessions.EnvironmentKey{TenantID: optionalID(row.TenantID), EnvironmentID: optionalID(row.ID)}) + } + return keys, nil +} + +// environmentTx is the Session's Environment inside the Session transaction +// of an initialization or connection operation. +type environmentTx struct{ *SessionTx } + +var ( + _ sessions.InitializationTx = (*environmentTx)(nil) + _ sessions.ConnectionTx = (*environmentTx)(nil) +) + +func (t *environmentTx) LoadSessionDevice(ctx context.Context) (sessions.ExecutionDevice, bool, error) { + return loadSessionDevice(ctx, t.q, t.tenant, t.session) +} + +func (t *environmentTx) ClaimInitialization(ctx context.Context, environment string) (bool, error) { + id, err := parseID(environment) + if err != nil { + return false, err + } + count, err := t.q.ClaimEnvironmentInitialization(ctx, id) + return count == 1, err +} + +func (t *environmentTx) CompleteInitialization(ctx context.Context, environment string) (bool, error) { + id, err := parseID(environment) + if err != nil { + return false, err + } + count, err := t.q.CompleteEnvironmentInitialization(ctx, id) + return count == 1, err +} + +func (t *environmentTx) FailInitialization(ctx context.Context, environment string) error { + id, err := parseID(environment) + if err != nil { + return err + } + return t.q.FailEnvironmentInitialization(ctx, id) +} + +func (t *environmentTx) LoadConnection(ctx context.Context, environment string) (sessions.EnvironmentConnection, bool, error) { + id, err := parseID(environment) + if err != nil { + return sessions.EnvironmentConnection{}, false, err + } + row, err := t.q.GetEnvironmentConnection(ctx, id) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.EnvironmentConnection{}, false, nil + } + if err != nil { + return sessions.EnvironmentConnection{}, false, err + } + return sessions.EnvironmentConnection{Generation: uuid.UUID(row.Generation.Bytes).String(), Revision: row.Revision}, true, nil +} + +func (t *environmentTx) ReplaceConnection(ctx context.Context, environment, generation string) error { + id, err := parseID(environment) + if err != nil { + return err + } + current, err := parseID(generation) + if err != nil { + return err + } + return t.q.ReplaceEnvironmentConnection(ctx, sqlc.ReplaceEnvironmentConnectionParams{EnvironmentID: id, Generation: current}) +} + +func (t *environmentTx) AdvanceConnection(ctx context.Context, environment string, revision int64) error { + id, err := parseID(environment) + if err != nil { + return err + } + return t.q.AdvanceEnvironmentConnection(ctx, sqlc.AdvanceEnvironmentConnectionParams{EnvironmentID: id, Revision: revision}) +} + +func (t *environmentTx) DeleteConnection(ctx context.Context, environment string) error { + id, err := parseID(environment) + if err != nil { + return err + } + return t.q.DeleteEnvironmentConnection(ctx, id) +} + +func (t *environmentTx) SetConnectionStatus(ctx context.Context, environment, status string) error { + id, err := parseID(environment) + if err != nil { + return err + } + return t.q.SetEnvironmentConnectionStatus(ctx, sqlc.SetEnvironmentConnectionStatusParams{ID: id, Status: status}) +} diff --git a/services/core/internal/persistence/postgres/sessionpg/executor_credentials.go b/services/core/internal/persistence/postgres/sessionpg/executor_credentials.go new file mode 100644 index 000000000..8238dd7e0 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/executor_credentials.go @@ -0,0 +1,35 @@ +package sessionpg + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + + "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/sessions" +) + +func (s *Store) AuthenticateEnvironmentExecutor(ctx context.Context, environment, digest string) (string, error) { + id, err := pgunit.ParseID(environment) + if err != nil { + return "", sessions.ErrNotFound + } + hash, err := hex.DecodeString(digest) + if err != nil || len(hash) != sha256.Size { + return "", sessions.ErrNotFound + } + tenant, err := s.units.Queries().AuthenticateEnvironmentExecutor(ctx, sqlc.AuthenticateEnvironmentExecutorParams{EnvironmentID: id, TokenSha256: hex.EncodeToString(hash)}) + if errors.Is(err, pgx.ErrNoRows) { + return "", sessions.ErrNotFound + } + if err != nil { + return "", fmt.Errorf("authenticate environment executor: %w", err) + } + return uuid.UUID(tenant.Bytes).String(), nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/frozen.go b/services/core/internal/persistence/postgres/sessionpg/frozen.go new file mode 100644 index 000000000..4b7e74cda --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/frozen.go @@ -0,0 +1,78 @@ +package sessionpg + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + + "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/environmentconfig" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// ReadEnvironmentSetup opens the setup the Session froze at creation. Frozen +// data that does not open, decode or validate is corrupt stored data, never +// the caller's input, so it is an internal error. +func (s *Store) ReadEnvironmentSetup(ctx context.Context, tenant, session string) (environmentconfig.Setup, error) { + lookup, err := DeviceLookup(tenant, session) + if err != nil { + return environmentconfig.Setup{}, sessions.ErrNotFound + } + encrypted, err := s.units.Queries().GetEnvironmentSetup(ctx, sqlc.GetEnvironmentSetupParams{TenantID: lookup.TenantID, ID: lookup.ID}) + if errors.Is(err, pgx.ErrNoRows) { + return environmentconfig.Setup{}, sessions.ErrNotFound + } + if err != nil { + return environmentconfig.Setup{}, err + } + var setup environmentconfig.Setup + if len(encrypted) == 0 { + return setup, nil + } + if s.cipher == nil { + return environmentconfig.Setup{}, credentialcrypto.ErrUnavailable + } + binding := credentialcrypto.EnvironmentSetupBinding{TenantID: uuid.UUID(lookup.TenantID.Bytes).String(), Resource: "session", OwnerID: uuid.UUID(lookup.ID.Bytes).String(), Field: "initialization"} + plaintext, err := s.cipher.OpenEnvironmentSetup(encrypted, binding) + if err != nil { + return environmentconfig.Setup{}, fmt.Errorf("open frozen environment setup: %w", err) + } + if environmentconfig.Decode(plaintext, &setup) != nil || setup.ValidateInstalled() != nil { + return environmentconfig.Setup{}, errors.New("frozen environment setup is invalid") + } + return setup, nil +} + +// ReadInitialEnvironmentFile opens the initial file the Session froze at +// position. A file that does not open or match its recorded size is corrupt +// stored data, so it is an internal error. +func (s *Store) ReadInitialEnvironmentFile(ctx context.Context, tenant, session string, position int) (environmentconfig.InitialFileMetadata, []byte, error) { + lookup, err := DeviceLookup(tenant, session) + if err != nil { + return environmentconfig.InitialFileMetadata{}, nil, err + } + row, err := s.units.Queries().GetInitialEnvironmentFile(ctx, sqlc.GetInitialEnvironmentFileParams{TenantID: lookup.TenantID, SessionID: lookup.ID, Position: int32(position)}) + if errors.Is(err, pgx.ErrNoRows) { + return environmentconfig.InitialFileMetadata{}, nil, sessions.ErrNotFound + } + if err != nil { + return environmentconfig.InitialFileMetadata{}, nil, err + } + if s.cipher == nil { + return environmentconfig.InitialFileMetadata{}, nil, credentialcrypto.ErrUnavailable + } + id := uuid.UUID(row.ID.Bytes).String() + binding := credentialcrypto.EnvironmentFileBinding{TenantID: uuid.UUID(lookup.TenantID.Bytes).String(), Resource: "session", OwnerID: uuid.UUID(lookup.ID.Bytes).String(), FileID: id} + body, err := s.cipher.OpenEnvironmentFile(row.Contents, binding) + if err != nil { + return environmentconfig.InitialFileMetadata{}, nil, fmt.Errorf("open frozen initial file: %w", err) + } + if int64(len(body)) != row.SizeBytes { + return environmentconfig.InitialFileMetadata{}, nil, errors.New("frozen initial file does not match its recorded size") + } + return environmentconfig.InitialFileMetadata{ID: id, Path: row.Path, SizeBytes: &row.SizeBytes}, body, nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/frozen_test.go b/services/core/internal/persistence/postgres/sessionpg/frozen_test.go new file mode 100644 index 000000000..da87dc7c9 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/frozen_test.go @@ -0,0 +1,141 @@ +package sessionpg + +import ( + "bytes" + "errors" + "reflect" + "testing" + + "github.com/google/uuid" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +func frozenCipher(t *testing.T) *credentialcrypto.Cipher { + t.Helper() + cipher, err := credentialcrypto.New(bytes.Repeat([]byte{5}, 32)) + if err != nil { + t.Fatal(err) + } + return cipher +} + +// corruptFrozenData reports whether err classifies frozen data as corrupt +// stored data: an internal error, never the caller's input, a missing record +// or a missing key. +func corruptFrozenData(err error) bool { + return err != nil && !errors.Is(err, sessions.ErrInvalidInput) && !errors.Is(err, sessions.ErrNotFound) && !errors.Is(err, credentialcrypto.ErrUnavailable) +} + +func TestReadEnvironmentSetupClassifiesFrozenData(t *testing.T) { + pool := pgtest.Open(t) + cipher := frozenCipher(t) + adapter := New(pgunit.NewPool(pool), cipher) + tenantID, sessionID, _ := newEnvironment(t, pool, "self_hosted", "pending") + tenant, session := uuidText(tenantID), uuidText(sessionID) + binding := credentialcrypto.EnvironmentSetupBinding{TenantID: tenant, Resource: "session", OwnerID: session, Field: "initialization"} + freeze := func(plaintext string, binding credentialcrypto.EnvironmentSetupBinding) []byte { + t.Helper() + sealed, err := cipher.SealEnvironmentSetup([]byte(plaintext), binding) + if err != nil { + t.Fatal(err) + } + return sealed + } + + if setup, err := adapter.ReadEnvironmentSetup(t.Context(), tenant, session); err != nil || !reflect.DeepEqual(setup, environmentconfig.Setup{}) { + t.Fatalf("Session without a frozen setup: %+v, %v", setup, err) + } + for name, ids := range map[string][2]string{"unknown Session": {tenant, uuid.NewString()}, "other tenant": {uuid.NewString(), session}, "malformed Session": {tenant, "session"}} { + if _, err := adapter.ReadEnvironmentSetup(t.Context(), ids[0], ids[1]); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("%s: %v", name, err) + } + } + + exec(t, pool, `INSERT INTO environment_setups(session_id, contents) VALUES ($1, $2)`, sessionID, freeze(`{"env":{"GREETING":"hello"},"packages":{"npm":["left-pad"],"python":[]}}`, binding)) + setup, err := adapter.ReadEnvironmentSetup(t.Context(), tenant, session) + if err != nil || setup.Env["GREETING"] != "hello" || !reflect.DeepEqual(setup.Packages.NPM, []string{"left-pad"}) { + t.Fatalf("frozen setup %+v, %v", setup, err) + } + if _, err := New(pgunit.NewPool(pool), nil).ReadEnvironmentSetup(t.Context(), tenant, session); !errors.Is(err, credentialcrypto.ErrUnavailable) { + t.Fatalf("without a credential key: %v", err) + } + + otherSession := binding + otherSession.OwnerID = uuid.NewString() + for name, contents := range map[string][]byte{ + "sealed for another Session": freeze(`{"packages":{"npm":[],"python":[]}}`, otherSession), + "not ciphertext": []byte("plaintext"), + "not JSON": freeze(`packages`, binding), + "unknown field": freeze(`{"packages":{"npm":[],"python":[]},"network":{}}`, binding), + "invalid setup": freeze(`{"env":{"PATH":"/bin"},"packages":{"npm":[],"python":[]}}`, binding), + "null system packages": freeze(`{"packages":{"npm":[],"python":[],"system":null}}`, binding), + "empty system packages": freeze(`{"packages":{"npm":[],"python":[],"system":[]}}`, binding), + "system packages": freeze(`{"packages":{"npm":[],"python":[],"system":["jq"]}}`, binding), + } { + exec(t, pool, `UPDATE environment_setups SET contents = $2 WHERE session_id = $1`, sessionID, contents) + if _, err := adapter.ReadEnvironmentSetup(t.Context(), tenant, session); !corruptFrozenData(err) { + t.Fatalf("%s: %v", name, err) + } + } + + exec(t, pool, `UPDATE sessions SET deleted_at = clock_timestamp() WHERE id = $1`, sessionID) + if _, err := adapter.ReadEnvironmentSetup(t.Context(), tenant, session); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("deleted Session: %v", err) + } +} + +func TestReadInitialEnvironmentFileClassifiesFrozenData(t *testing.T) { + pool := pgtest.Open(t) + cipher := frozenCipher(t) + adapter := New(pgunit.NewPool(pool), cipher) + tenantID, sessionID, _ := newEnvironment(t, pool, "self_hosted", "pending") + tenant, session := uuidText(tenantID), uuidText(sessionID) + file := uuid.NewString() + binding := credentialcrypto.EnvironmentFileBinding{TenantID: tenant, Resource: "session", OwnerID: session, FileID: file} + freeze := func(body string, binding credentialcrypto.EnvironmentFileBinding) []byte { + t.Helper() + sealed, err := cipher.SealEnvironmentFile([]byte(body), binding) + if err != nil { + t.Fatal(err) + } + return sealed + } + exec(t, pool, `INSERT INTO initial_environment_files(id, session_id, position, path, size_bytes, contents) VALUES ($1, $2, 0, 'notes.txt', 5, $3)`, file, sessionID, freeze("hello", binding)) + + metadata, body, err := adapter.ReadInitialEnvironmentFile(t.Context(), tenant, session, 0) + if err != nil || metadata.ID != file || metadata.Path != "notes.txt" || metadata.SizeBytes == nil || *metadata.SizeBytes != 5 || string(body) != "hello" { + t.Fatalf("frozen file %+v %q, %v", metadata, body, err) + } + if _, _, err := adapter.ReadInitialEnvironmentFile(t.Context(), tenant, session, 1); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("missing position: %v", err) + } + if _, _, err := adapter.ReadInitialEnvironmentFile(t.Context(), uuid.NewString(), session, 0); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("other tenant: %v", err) + } + if _, _, err := New(pgunit.NewPool(pool), nil).ReadInitialEnvironmentFile(t.Context(), tenant, session, 0); !errors.Is(err, credentialcrypto.ErrUnavailable) { + t.Fatalf("without a credential key: %v", err) + } + + otherFile := binding + otherFile.FileID = uuid.NewString() + for name, contents := range map[string][]byte{ + "sealed for another file": freeze("hello", otherFile), + "not ciphertext": []byte("hello"), + "size mismatch": freeze("hello!", binding), + } { + exec(t, pool, `UPDATE initial_environment_files SET contents = $2 WHERE id = $1`, file, contents) + if _, _, err := adapter.ReadInitialEnvironmentFile(t.Context(), tenant, session, 0); !corruptFrozenData(err) { + t.Fatalf("%s: %v", name, err) + } + } + + exec(t, pool, `UPDATE sessions SET deleted_at = clock_timestamp() WHERE id = $1`, sessionID) + if _, _, err := adapter.ReadInitialEnvironmentFile(t.Context(), tenant, session, 0); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("deleted Session: %v", err) + } +} diff --git a/services/core/internal/persistence/postgres/sessionpg/session.go b/services/core/internal/persistence/postgres/sessionpg/session.go index 162dabe26..2daebf41a 100644 --- a/services/core/internal/persistence/postgres/sessionpg/session.go +++ b/services/core/internal/persistence/postgres/sessionpg/session.go @@ -36,14 +36,21 @@ func LockSession(ctx context.Context, q *sqlc.Queries, tenant, session pgtype.UU // with LockSession. func WithSession(ctx context.Context, runner pgunit.Transactor, tenant, session pgtype.UUID, apply func(context.Context, *sqlc.Queries, sessions.LockedSession) error) error { return runner.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := sqlc.New(tx) - locked, err := LockSession(ctx, q, tenant, session) - if err != nil { - return err - } - if err := apply(ctx, q, locked); err != nil { - return err - } - return PruneChanges(ctx, q, session) + return inSession(ctx, sqlc.New(tx), tenant, session, apply) }) } + +// inSession runs the body of a Session transaction on its queries: it locks +// the tenant's Session, runs apply with what the lock shows, then prunes the +// Session's journal. Session transactions that must read before they know the +// Session, such as one keyed by its Environment, run it after that read. +func inSession(ctx context.Context, q *sqlc.Queries, tenant, session pgtype.UUID, apply func(context.Context, *sqlc.Queries, sessions.LockedSession) error) error { + locked, err := LockSession(ctx, q, tenant, session) + if err != nil { + return err + } + if err := apply(ctx, q, locked); err != nil { + return err + } + return PruneChanges(ctx, q, session) +} diff --git a/services/core/internal/persistence/postgres/sessionpg/store.go b/services/core/internal/persistence/postgres/sessionpg/store.go index fb32417b5..f595a647f 100644 --- a/services/core/internal/persistence/postgres/sessionpg/store.go +++ b/services/core/internal/persistence/postgres/sessionpg/store.go @@ -1,16 +1,24 @@ package sessionpg import ( + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) // Store is the pooled Session adapter: the storage of the Session use cases -// and the Session reads. -type Store struct{ units *pgunit.Pool } +// and the Session reads. cipher opens the Session data frozen at creation; it +// is nil on a service without a credential key. +type Store struct { + units *pgunit.Pool + cipher *credentialcrypto.Cipher +} -// New builds the pooled Session adapter on units. -func New(units *pgunit.Pool) *Store { return &Store{units: units} } +// New builds the pooled Session adapter on units, opening frozen Session data +// with cipher. +func New(units *pgunit.Pool, cipher *credentialcrypto.Cipher) *Store { + return &Store{units: units, cipher: cipher} +} var ( _ sessions.Storage = (*Store)(nil) diff --git a/services/core/internal/runtime/gateway.go b/services/core/internal/runtime/gateway.go index 84d131ae8..56a4a06e1 100644 --- a/services/core/internal/runtime/gateway.go +++ b/services/core/internal/runtime/gateway.go @@ -11,22 +11,20 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" ) -type DeviceStore interface { - runtimegateway.RuntimeStore - runtimegateway.HeartbeatTouch -} - // NewGateway serves the V1 daemon executor transport for both managed and -// user-managed Runtime. Its credentials never grant public Session API access. -func NewGateway(s DeviceStore, publicWSURL string) (http.Handler, *runtimegateway.Registry, error) { +// user-managed Runtime. It authenticates devices with credentials, records +// their heartbeats with heartbeat and drains an archived Session's cancellation +// receipts through cancellations. Its credentials never grant public Session +// API access. +func NewGateway(credentials runtimegateway.RuntimeStore, heartbeat runtimegateway.HeartbeatTouch, cancellations runtimegateway.ArchivedCancellationStore, publicWSURL string) (http.Handler, *runtimegateway.Registry, error) { u, err := url.Parse(publicWSURL) - if err != nil || s == nil || (u.Scheme != "ws" && u.Scheme != "wss") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" || u.Path != "/api/v1/agent-daemon/ws" { + if err != nil || credentials == nil || heartbeat == nil || cancellations == nil || (u.Scheme != "ws" && u.Scheme != "wss") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" || u.Path != "/api/v1/agent-daemon/ws" { return nil, nil, errors.New("daemon URL must be an absolute ws(s) URL ending in /api/v1/agent-daemon/ws") } registry := runtimegateway.NewRegistry() h := runtimegateway.NewHandler(runtimegateway.HandlerConfig{ - Authenticator: runtimegateway.NewAuthenticator(s), Registry: registry, - Heartbeat: s, PublicWSURL: publicWSURL, + Authenticator: runtimegateway.NewAuthenticator(credentials), Registry: registry, + Heartbeat: heartbeat, ArchivedCancellations: cancellations, PublicWSURL: publicWSURL, }) r := chi.NewRouter() r.Route("/api/v1", func(r chi.Router) { runtimegateway.RegisterRoutes(r, h) }) diff --git a/services/core/internal/runtimedevice/credential.go b/services/core/internal/runtimedevice/credential.go index daede9156..b108f00d8 100644 --- a/services/core/internal/runtimedevice/credential.go +++ b/services/core/internal/runtimedevice/credential.go @@ -6,6 +6,10 @@ import ( "strings" ) +// RuntimeTypeAgentDaemon is the type of every device credential: the device +// runs the agent daemon. The gateway authenticates only this type. +const RuntimeTypeAgentDaemon = "agent_daemon" + // Credential is the minimal device identity needed for gateway authentication. // CredentialHash is never sent over the daemon protocol. type Credential struct { diff --git a/services/core/internal/runtimeenrollment/connection_test.go b/services/core/internal/runtimeenrollment/connection_test.go index 23fd1aff2..e10519f92 100644 --- a/services/core/internal/runtimeenrollment/connection_test.go +++ b/services/core/internal/runtimeenrollment/connection_test.go @@ -108,7 +108,7 @@ func (s *liveConnectionStore) GetDeviceCredential(context.Context, string) (runt if s.deviceRevokedAtRecheck && s.authCalls >= 2 { return runtimedevice.Credential{}, false, nil } - return runtimedevice.Credential{ID: "device", WorkspaceID: "tenant", Type: runtimegateway.RuntimeTypeAgentDaemon, CredentialHash: s.digest}, true, nil + return runtimedevice.Credential{ID: "device", WorkspaceID: "tenant", Type: runtimedevice.RuntimeTypeAgentDaemon, CredentialHash: s.digest}, true, nil } func TestRuntimeConnectedCurrentAuthorityAfterPeer(t *testing.T) { for _, name := range []string{"connected", "rotated before read", "revoked after peer", "device revoked after peer", "retired after peer", "store error after peer", "device store error after peer", "closed", "no registry"} { diff --git a/services/core/internal/runtimeenrollment/enrollment.go b/services/core/internal/runtimeenrollment/enrollment.go index d2f0d2e5c..2a2ff3ce6 100644 --- a/services/core/internal/runtimeenrollment/enrollment.go +++ b/services/core/internal/runtimeenrollment/enrollment.go @@ -11,11 +11,10 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) type EnrollmentStore interface { - EnrollRuntime(context.Context, string, string) (store.RuntimeEnrollment, error) + EnrollRuntime(context.Context, string, string) (sessions.RuntimeEnrollment, error) } // EnrollmentHandler is part of our daemon connection contract, not an upstream diff --git a/services/core/internal/runtimeenrollment/enrollment_test.go b/services/core/internal/runtimeenrollment/enrollment_test.go index 42359248c..6eb774833 100644 --- a/services/core/internal/runtimeenrollment/enrollment_test.go +++ b/services/core/internal/runtimeenrollment/enrollment_test.go @@ -9,7 +9,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) type enrollmentStub struct { @@ -17,12 +16,12 @@ type enrollmentStub struct { err error } -func (s *enrollmentStub) EnrollRuntime(_ context.Context, environment, digest string) (store.RuntimeEnrollment, error) { +func (s *enrollmentStub) EnrollRuntime(_ context.Context, environment, digest string) (sessions.RuntimeEnrollment, error) { s.calls++ if environment != "environment" || digest != runtimedevice.HashCredential("private-test-token") { - return store.RuntimeEnrollment{}, errors.New("unexpected enrollment input") + return sessions.RuntimeEnrollment{}, errors.New("unexpected enrollment input") } - return store.RuntimeEnrollment{DeviceID: "device", SessionID: "session", EnvironmentID: environment, WorkspaceDirectory: "/workspace"}, s.err + return sessions.RuntimeEnrollment{DeviceID: "device", SessionID: "session", EnvironmentID: environment, WorkspaceDirectory: "/workspace"}, s.err } func TestEnrollmentConnectionContract(t *testing.T) { diff --git a/services/core/internal/runtimegateway/auth.go b/services/core/internal/runtimegateway/auth.go index e20bed650..4d7fadea2 100644 --- a/services/core/internal/runtimegateway/auth.go +++ b/services/core/internal/runtimegateway/auth.go @@ -9,12 +9,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" ) -// RuntimeTypeAgentDaemon is the value runtimes.type takes for rows -// that back an agent_daemon device. The DB has no CHECK constraint on -// runtimes.type so this is a Go-side invariant; the listRuntimes admin -// endpoint's allowlist must be updated alongside it. -const RuntimeTypeAgentDaemon = "agent_daemon" - var ErrAuthMissingParams = errors.New("agentdaemon auth: missing device_id / token / version") var ErrAuthUnknownDevice = errors.New("agentdaemon auth: unknown device") @@ -72,7 +66,7 @@ func (a *Authenticator) AuthenticateBearer(ctx context.Context, deviceID, bearer if !ok { return AuthenticatedRuntime{}, ErrAuthUnknownDevice } - if rt.Type != RuntimeTypeAgentDaemon { + if rt.Type != runtimedevice.RuntimeTypeAgentDaemon { return AuthenticatedRuntime{}, ErrAuthWrongRuntimeType } storedHash := rt.CredentialHash diff --git a/services/core/internal/runtimegateway/auth_test.go b/services/core/internal/runtimegateway/auth_test.go index a480baf43..4aeb25182 100644 --- a/services/core/internal/runtimegateway/auth_test.go +++ b/services/core/internal/runtimegateway/auth_test.go @@ -57,7 +57,7 @@ func TestAuthenticator_RejectsWrongRuntimeType(t *testing.T) { func TestAuthenticator_RejectsBadCredential(t *testing.T) { row := runtimedevice.Credential{ ID: "dev-1", - Type: RuntimeTypeAgentDaemon, + Type: runtimedevice.RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential("real-tok"), } auth := NewAuthenticator(&stubRuntimeStore{row: row, ok: true}) @@ -82,7 +82,7 @@ func TestAuthenticator_AcceptsValidCredential(t *testing.T) { ID: "dev-1", WorkspaceID: "wks-1", Name: "alice-mac", - Type: RuntimeTypeAgentDaemon, + Type: runtimedevice.RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential("real-tok"), } auth := NewAuthenticator(&stubRuntimeStore{row: row, ok: true}) diff --git a/services/core/internal/runtimegateway/bootstrap_url_test.go b/services/core/internal/runtimegateway/bootstrap_url_test.go index 70f379122..7b32c8d82 100644 --- a/services/core/internal/runtimegateway/bootstrap_url_test.go +++ b/services/core/internal/runtimegateway/bootstrap_url_test.go @@ -13,7 +13,7 @@ func TestBootstrapPublishesURLOnlyAfterAuthentication(t *testing.T) { const publicURL = "wss://core.example/api/v1/agent-daemon/ws" for _, valid := range []bool{true, false} { h := NewHandler(HandlerConfig{Registry: NewRegistry(), PublicWSURL: publicURL, - Authenticator: NewAuthenticator(&stubRuntimeStore{ok: true, row: runtimedevice.Credential{ID: "device", WorkspaceID: "tenant", Type: RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential(token)}})}) + Authenticator: NewAuthenticator(&stubRuntimeStore{ok: true, row: runtimedevice.Credential{ID: "device", WorkspaceID: "tenant", Type: runtimedevice.RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential(token)}})}) r := httptest.NewRequest(http.MethodPost, "/agent-daemon/bootstrap", strings.NewReader(`{"device_id":"device"}`)) bearer, status := "wrong", http.StatusUnauthorized if valid { diff --git a/services/core/internal/runtimegateway/cancellation.go b/services/core/internal/runtimegateway/cancellation.go index cf1f2f6d5..2ee2b13ed 100644 --- a/services/core/internal/runtimegateway/cancellation.go +++ b/services/core/internal/runtimegateway/cancellation.go @@ -49,8 +49,7 @@ func (s *Session) TrackExecutionDelivery(runID string) (func(), error) { // DrainArchivedCancellation checks both durable archive identity and this exact // connection's unfinished delivery. It performs no liveness or credential write. func (s *Session) DrainArchivedCancellation(ctx context.Context) (bool, error) { - store, ok := s.heartbeat.(ArchivedCancellationStore) - if !ok || s.IsClosed() { + if s.archivedCancellations == nil || s.IsClosed() { return false, nil } s.receiptMu.Lock() @@ -62,7 +61,7 @@ func (s *Session) DrainArchivedCancellation(ctx context.Context) (bool, error) { if len(ids) == 0 { return false, nil } - receipt, err := store.ArchivedCancellationReceipt(ctx, s.DeviceID, s.credentialHash, ids) + receipt, err := s.archivedCancellations.ArchivedCancellationReceipt(ctx, s.DeviceID, s.credentialHash, ids) if err != nil || receipt.RunID == "" { return false, err } diff --git a/services/core/internal/runtimegateway/cancellation_test.go b/services/core/internal/runtimegateway/cancellation_test.go index a8e099f9d..aeae810b7 100644 --- a/services/core/internal/runtimegateway/cancellation_test.go +++ b/services/core/internal/runtimegateway/cancellation_test.go @@ -11,7 +11,6 @@ import ( ) type receiptStore struct { - *fakeHeartbeatStore mu sync.Mutex receipt runtimedevice.ArchivedCancellationReceipt } @@ -36,7 +35,8 @@ func TestArchivedReceiptTracksDeliveryBeyondDoneAndRejectsNewWork(t *testing.T) defer peer.Close("test complete") peer.credentialHash = "original-hash" receipt := runtimedevice.ArchivedCancellationReceipt{RunID: "run", Deadline: time.Now().Add(time.Second)} - peer.heartbeat = &receiptStore{fakeHeartbeatStore: newFakeHeartbeatStore(), receipt: receipt} + peer.heartbeat = newFakeHeartbeatStore() + peer.archivedCancellations = &receiptStore{receipt: receipt} if draining, err := peer.DrainArchivedCancellation(t.Context()); err != nil || draining { t.Fatal("unowned delivery got drain", draining, err) } @@ -127,7 +127,8 @@ func TestArchivedReceiptDeadlineDoesNotRenew(t *testing.T) { defer peer.Close("test complete") peer.credentialHash = "original-hash" receipt := runtimedevice.ArchivedCancellationReceipt{RunID: "run", Deadline: time.Now().Add(100 * time.Millisecond)} - peer.heartbeat = &receiptStore{fakeHeartbeatStore: newFakeHeartbeatStore(), receipt: receipt} + peer.heartbeat = newFakeHeartbeatStore() + peer.archivedCancellations = &receiptStore{receipt: receipt} release, err := peer.TrackExecutionDelivery("run") if err != nil { t.Fatal(err) diff --git a/services/core/internal/runtimegateway/handler.go b/services/core/internal/runtimegateway/handler.go index a34782d40..85a62d55e 100644 --- a/services/core/internal/runtimegateway/handler.go +++ b/services/core/internal/runtimegateway/handler.go @@ -20,7 +20,6 @@ import ( type HeartbeatTouch interface { TouchRuntimeHeartbeat(ctx context.Context, runtimeID string) (runtimedevice.HeartbeatStatus, error) TouchAgentDaemonHeartbeat(ctx context.Context, input runtimedevice.Heartbeat) (runtimedevice.HeartbeatStatus, error) - MarkRuntimeOffline(ctx context.Context, runtimeID string) error } // HandlerConfig wires the gateway's HTTP/WS handlers. nil values panic @@ -37,6 +36,10 @@ type HandlerConfig struct { // last_heartbeat_at fresh. nil tracks liveness in-process only. Heartbeat HeartbeatTouch + // ArchivedCancellations reads the receipt that an archived Session's + // cancellation still owes a connection's delivery. nil drains nothing. + ArchivedCancellations ArchivedCancellationStore + // PublicWSURL is the wss://... URL returned in the bootstrap // response so deployments behind a TLS terminator can advertise // the externally-reachable URL. @@ -164,6 +167,7 @@ func (h *Handler) WS(w http.ResponseWriter, r *http.Request) { } sess := NewSessionWithOwner(conn, auth.DeviceID, auth.WorkspaceID, version, h.cfg.Registry, h.cfg.Log, lease) sess.heartbeat = h.cfg.Heartbeat + sess.archivedCancellations = h.cfg.ArchivedCancellations sess.credentialHash = runtimedevice.HashCredential(token) h.cfg.Log("agentdaemon gateway: ws upgrade ok, registering device_id=%s owner_pod=%s waiters=%d", auth.DeviceID, h.cfg.OwnerPodID, len(h.cfg.Registry.PendingWaiters(auth.DeviceID))) diff --git a/services/core/internal/runtimegateway/handler_test.go b/services/core/internal/runtimegateway/handler_test.go index e2e03cd63..59b853cc8 100644 --- a/services/core/internal/runtimegateway/handler_test.go +++ b/services/core/internal/runtimegateway/handler_test.go @@ -31,7 +31,7 @@ func TestWebSocketRequiresAuthorizationBearer(t *testing.T) { handler := NewHandler(HandlerConfig{ Registry: registry, Authenticator: NewAuthenticator(&stubRuntimeStore{ok: true, row: runtimedevice.Credential{ - ID: "device", WorkspaceID: "tenant", Type: RuntimeTypeAgentDaemon, + ID: "device", WorkspaceID: "tenant", Type: runtimedevice.RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential(credential), }}), }) diff --git a/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go b/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go index 778e4032d..7a6982fca 100644 --- a/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go +++ b/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go @@ -51,7 +51,7 @@ func TestLiveMCPBearerGatewayColdContinuation(t *testing.T) { }) id := uuid.NewString() registry := NewRegistry() - auth := NewAuthenticator(&stubRuntimeStore{ok: true, row: runtimedevice.Credential{ID: id, WorkspaceID: uuid.NewString(), Type: RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential(runner)}}) + auth := NewAuthenticator(&stubRuntimeStore{ok: true, row: runtimedevice.Credential{ID: id, WorkspaceID: uuid.NewString(), Type: runtimedevice.RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential(runner)}}) router := chi.NewRouter() server := httptest.NewServer(router) t.Cleanup(server.Close) diff --git a/services/core/internal/runtimegateway/session.go b/services/core/internal/runtimegateway/session.go index d2fd77d35..cb115fedb 100644 --- a/services/core/internal/runtimegateway/session.go +++ b/services/core/internal/runtimegateway/session.go @@ -90,6 +90,9 @@ type Session struct { // heartbeat persists daemon-advertised capability snapshots. heartbeat HeartbeatTouch credentialHash string + // archivedCancellations reads the receipt an archived Session's + // cancellation owes this connection's delivery. + archivedCancellations ArchivedCancellationStore hbMu sync.Mutex lastSeenAt time.Time @@ -258,7 +261,6 @@ func (s *Session) Close(reason string) { s.closeWorkspaceWrites() s.closeCapabilities() s.closeWorkspaceExports() - s.markOfflineOnClose() s.releaseOwnerLease() }) } @@ -451,17 +453,6 @@ func (s *Session) releaseOwnerLease() { } } -func (s *Session) markOfflineOnClose() { - if s.heartbeat == nil { - return - } - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) - defer cancel() - if err := s.heartbeat.MarkRuntimeOffline(ctx, s.DeviceID); err != nil { - s.log("agentdaemon gateway: mark offline on close failed device=%s: %v", s.DeviceID, err) - } -} - func (s *Session) handleHeartbeat(env proto.Envelope) { var p proto.HeartbeatPayload if err := env.DecodePayload(&p); err != nil { diff --git a/services/core/internal/runtimegateway/session_test.go b/services/core/internal/runtimegateway/session_test.go index fc1422ccf..a30e8a66d 100644 --- a/services/core/internal/runtimegateway/session_test.go +++ b/services/core/internal/runtimegateway/session_test.go @@ -132,10 +132,6 @@ func (f *fakeHeartbeatStore) TouchAgentDaemonHeartbeat(_ context.Context, input return runtimedevice.HeartbeatStatus{Liveness: "online"}, nil } -func (f *fakeHeartbeatStore) MarkRuntimeOffline(_ context.Context, _ string) error { - return nil -} - func (f *fakeHeartbeatStore) waitDaemonHeartbeat(t *testing.T) runtimedevice.Heartbeat { t.Helper() select { diff --git a/services/core/internal/runtimegateway/wire_test.go b/services/core/internal/runtimegateway/wire_test.go index a3b4acb64..7b09c5e9b 100644 --- a/services/core/internal/runtimegateway/wire_test.go +++ b/services/core/internal/runtimegateway/wire_test.go @@ -28,7 +28,7 @@ const wait = 3 * time.Second type credentialStore struct{} func (credentialStore) GetDeviceCredential(context.Context, string) (runtimedevice.Credential, bool, error) { - return runtimedevice.Credential{ID: prototest.DeviceID, WorkspaceID: "tenant", Type: runtimegateway.RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential(prototest.Credential)}, true, nil + return runtimedevice.Credential{ID: prototest.DeviceID, WorkspaceID: "tenant", Type: runtimedevice.RuntimeTypeAgentDaemon, CredentialHash: runtimedevice.HashCredential(prototest.Credential)}, true, nil } func newGateway(t *testing.T) (*runtimegateway.Registry, string) { diff --git a/services/core/internal/sandbox/providers/configuration_flow_test.go b/services/core/internal/sandbox/providers/configuration_flow_test.go index a37977ccf..88d4f46d3 100644 --- a/services/core/internal/sandbox/providers/configuration_flow_test.go +++ b/services/core/internal/sandbox/providers/configuration_flow_test.go @@ -137,7 +137,7 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { Artifacts: struct{ api.Artifacts }{}, ArtifactsReader: struct{ api.ArtifactsReader }{}, SessionAdmin: s, - Environments: s, Admin: s, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, + Environments: s, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, Admin: s, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, ExecutorConnections: struct{ api.ExecutorConnections }{}, Metrics: struct{ api.Metrics }{}, RuntimeObservations: struct{ api.RuntimeObservations }{}, RuntimeHistory: struct{ api.RuntimeHistory }{}, Execution: &api.Execution{ diff --git a/services/core/internal/sessions/artifacts_test.go b/services/core/internal/sessions/artifacts_test.go index d36891a3d..0b2b4561c 100644 --- a/services/core/internal/sessions/artifacts_test.go +++ b/services/core/internal/sessions/artifacts_test.go @@ -190,6 +190,31 @@ func (s *fakeArtifactStorage) DeleteSessionArtifact(context.Context, string, str return nil } +func (s *fakeArtifactStorage) CreateDevice(context.Context, string, DeviceRegistration) (ExecutionDevice, error) { + s.t.Fatal("unexpected call to CreateDevice") + return ExecutionDevice{}, nil +} + +func (s *fakeArtifactStorage) RevokeDevice(context.Context, string, string) error { + s.t.Fatal("unexpected call to RevokeDevice") + return nil +} + +func (s *fakeArtifactStorage) TouchDevice(context.Context, string) (bool, error) { + s.t.Fatal("unexpected call to TouchDevice") + return false, nil +} + +func (s *fakeArtifactStorage) TouchAuthenticatedDevice(context.Context, string, string) (bool, error) { + s.t.Fatal("unexpected call to TouchAuthenticatedDevice") + return false, nil +} + +func (s *fakeArtifactStorage) WithEnrollment(context.Context, string, string, func(context.Context, EnrollmentTx, Environment, LockedSession) error) error { + s.t.Fatal("unexpected call to WithEnrollment") + return nil +} + func discard(content io.Reader, size int64) error { _, err := io.CopyN(io.Discard, content, size) return err diff --git a/services/core/internal/sessions/devices.go b/services/core/internal/sessions/devices.go new file mode 100644 index 000000000..b419c3983 --- /dev/null +++ b/services/core/internal/sessions/devices.go @@ -0,0 +1,177 @@ +package sessions + +import ( + "context" + "encoding/hex" + "fmt" + "strings" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" +) + +// DeviceRegistration is a Runtime device's name and the SHA-256 digest of its +// credential, as NewDeviceRegistration validates them. +type DeviceRegistration struct { + Name string + CredentialHash string +} + +// NewDeviceRegistration validates a device registration: a name of 1 to 256 +// bytes once trimmed, and a SHA-256 credential digest in hex, which it +// normalizes to lowercase. Anything else is ErrInvalidInput. +func NewDeviceRegistration(name, credentialHash string) (DeviceRegistration, error) { + name = strings.TrimSpace(name) + digest, err := hex.DecodeString(credentialHash) + if err != nil || len(digest) != 32 || name == "" || len(name) > 256 { + return DeviceRegistration{}, fmt.Errorf("%w: device name and SHA-256 credential digest required", ErrInvalidInput) + } + return DeviceRegistration{Name: name, CredentialHash: hex.EncodeToString(digest)}, nil +} + +// RuntimeEnrollment is the resource binding an enrolled user-managed Runtime +// receives. It never carries another secret. +type RuntimeEnrollment struct { + DeviceID string + SessionID string + EnvironmentID string + WorkspaceDirectory string +} + +// EnrolledRuntimeBinding identifies user-managed compute, without an +// allocation or any promise of live authorization. The Worker rechecks the +// socket's key. +type EnrolledRuntimeBinding struct { + DeviceID, TenantID, EnvironmentID, SessionID string +} + +// EnrollmentAuthority is what an executor credential authorizes when it +// enrolls a Runtime: the key and the Environment's workspace directory. +type EnrollmentAuthority struct { + KeyID string + WorkspaceDirectory string +} + +// DeviceReader reads Runtime devices and their Session bindings. +type DeviceReader interface { + // GetSessionDevice reads the Runtime device bound to the Session once its + // Environment preparation completed; before that, and without a bound + // device, it is ErrNotFound. + GetSessionDevice(ctx context.Context, tenant, session string) (ExecutionDevice, error) + // GetSessionRuntimeDevice reads the Runtime device bound to the Session, + // or ErrNotFound. It reports an authorized connection binding and does not + // admit native execution or file access before preparation completes. + GetSessionRuntimeDevice(ctx context.Context, tenant, session string) (ExecutionDevice, error) + // GetDeviceCredential reads the credential the Runtime gateway + // authenticates a device with, and reports whether the device still has + // authority. + GetDeviceCredential(ctx context.Context, device string) (runtimedevice.Credential, bool, error) + // ListEnrolledRuntimeBindings lists the enrolled user-managed Runtimes of + // live Environments. + ListEnrolledRuntimeBindings(ctx context.Context) ([]EnrolledRuntimeBinding, error) +} + +// DeviceStorage stores Runtime devices. +type DeviceStorage interface { + // CreateDevice stores a new device of the tenant and returns it. + CreateDevice(ctx context.Context, tenant string, registration DeviceRegistration) (ExecutionDevice, error) + // RevokeDevice revokes the tenant's device; an unknown device is + // ErrNotFound. + RevokeDevice(ctx context.Context, tenant, device string) error + // TouchDevice records that the device was seen and reports whether it + // still has authority. + TouchDevice(ctx context.Context, device string) (bool, error) + // TouchAuthenticatedDevice records that the device was seen with the + // credential and reports whether that credential still has authority. + TouchAuthenticatedDevice(ctx context.Context, device, credentialHash string) (bool, error) + // WithEnrollment authenticates the executor credential for the + // Environment, then runs apply in the transaction of the Environment's + // Session, with the Environment and what the Session lock shows. A + // credential that does not authenticate is ErrNotFound. + WithEnrollment(ctx context.Context, environment, credentialHash string, apply func(context.Context, EnrollmentTx, Environment, LockedSession) error) error +} + +// EnrollmentTx is the Session transaction EnrollRuntime runs in. +type EnrollmentTx interface { + // AuthorizeEnrollment rechecks, under the Session lock, that the + // credential still authorizes enrolling a Runtime for the live self_hosted + // Environment, and holds the key's lock until the transaction ends, so + // revocation cannot race enrollment. Without that authority it is + // ErrNotFound. + AuthorizeEnrollment(ctx context.Context) (EnrollmentAuthority, error) + // EnrollDevice creates the Environment's user-managed Runtime device for + // the key, or returns the device the same key already enrolled. A device + // of another key, or a revoked one, is ErrDeviceBindingConflict. + EnrollDevice(ctx context.Context, key string) (string, error) + // BindDevice binds the device to the Session. A Session bound to another + // device is ErrDeviceBindingConflict. + BindDevice(ctx context.Context, device string) error +} + +// CreateDevice provisions a device for an operator. It is not a tenant-facing +// registration API. +func (s *Service) CreateDevice(ctx context.Context, tenant, name, credentialHash string) (ExecutionDevice, error) { + registration, err := NewDeviceRegistration(name, credentialHash) + if err != nil { + return ExecutionDevice{}, err + } + return s.storage.CreateDevice(ctx, tenant, registration) +} + +// RevokeDevice revokes the tenant's device. +func (s *Service) RevokeDevice(ctx context.Context, tenant, device string) error { + return s.storage.RevokeDevice(ctx, tenant, device) +} + +// TouchRuntimeHeartbeat records a Runtime connection. A device that lost its +// authority reports Deleted. +func (s *Service) TouchRuntimeHeartbeat(ctx context.Context, device string) (runtimedevice.HeartbeatStatus, error) { + current, err := s.storage.TouchDevice(ctx, device) + if err != nil { + return runtimedevice.HeartbeatStatus{}, err + } + return heartbeatStatus(current), nil +} + +// TouchAgentDaemonHeartbeat records an agent daemon heartbeat under the +// credential the gateway authenticated. A credential that lost its authority +// reports Deleted. +func (s *Service) TouchAgentDaemonHeartbeat(ctx context.Context, heartbeat runtimedevice.Heartbeat) (runtimedevice.HeartbeatStatus, error) { + current, err := s.storage.TouchAuthenticatedDevice(ctx, heartbeat.RuntimeID, heartbeat.CredentialHash) + if err != nil { + return runtimedevice.HeartbeatStatus{}, err + } + return heartbeatStatus(current), nil +} + +// heartbeatStatus is the liveness a heartbeat reports. Live connectivity +// belongs to the gateway Registry; only the last-seen time is stored. +func heartbeatStatus(current bool) runtimedevice.HeartbeatStatus { + return runtimedevice.HeartbeatStatus{Liveness: "online", Deleted: !current} +} + +// EnrollRuntime binds a user-managed Runtime to one self_hosted Environment +// with the executor credential whose digest is credentialHash. A retry keeps +// the same device and key; it cannot replace compute or adopt another native +// history. +func (s *Service) EnrollRuntime(ctx context.Context, environment, credentialHash string) (RuntimeEnrollment, error) { + var result RuntimeEnrollment + err := s.storage.WithEnrollment(ctx, environment, credentialHash, func(ctx context.Context, tx EnrollmentTx, current Environment, locked LockedSession) error { + if err := locked.Public(); err != nil { + return err + } + authority, err := tx.AuthorizeEnrollment(ctx) + if err != nil { + return err + } + device, err := tx.EnrollDevice(ctx, authority.KeyID) + if err != nil { + return err + } + if err := tx.BindDevice(ctx, device); err != nil { + return err + } + result = RuntimeEnrollment{DeviceID: device, SessionID: current.SessionID, EnvironmentID: current.ID, WorkspaceDirectory: authority.WorkspaceDirectory} + return nil + }) + return result, err +} diff --git a/services/core/internal/sessions/devices_test.go b/services/core/internal/sessions/devices_test.go new file mode 100644 index 000000000..00b8fdd64 --- /dev/null +++ b/services/core/internal/sessions/devices_test.go @@ -0,0 +1,194 @@ +package sessions + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" +) + +// fakeDeviceStorage is strict pooled Session storage for the device use +// cases. Each method records its call, with the arguments that identify it, +// then runs its func; a method whose func is unset fails the test. +type fakeDeviceStorage struct { + t *testing.T + calls []string + + createDevice func(DeviceRegistration) (ExecutionDevice, error) + revokeDevice func() error + touchDevice func() (bool, error) + touchAuthenticatedDevice func() (bool, error) + // enrollment is the transaction WithEnrollment applies in, with + // environment and locked. + enrollment *fakeTx + environment Environment + locked LockedSession +} + +func (s *fakeDeviceStorage) record(name string, set bool, detail ...string) { + s.t.Helper() + if !set { + s.t.Fatalf("unexpected call to %s", name) + } + s.calls = append(s.calls, strings.Join(append([]string{name}, detail...), " ")) +} + +func (s *fakeDeviceStorage) CreateDevice(_ context.Context, tenant string, registration DeviceRegistration) (ExecutionDevice, error) { + s.record("CreateDevice", s.createDevice != nil, tenant, registration.Name, registration.CredentialHash) + return s.createDevice(registration) +} + +func (s *fakeDeviceStorage) RevokeDevice(_ context.Context, tenant, device string) error { + s.record("RevokeDevice", s.revokeDevice != nil, tenant, device) + return s.revokeDevice() +} + +func (s *fakeDeviceStorage) TouchDevice(_ context.Context, device string) (bool, error) { + s.record("TouchDevice", s.touchDevice != nil, device) + return s.touchDevice() +} + +func (s *fakeDeviceStorage) TouchAuthenticatedDevice(_ context.Context, device, credentialHash string) (bool, error) { + s.record("TouchAuthenticatedDevice", s.touchAuthenticatedDevice != nil, device, credentialHash) + return s.touchAuthenticatedDevice() +} + +func (s *fakeDeviceStorage) WithEnrollment(ctx context.Context, environment, credentialHash string, apply func(context.Context, EnrollmentTx, Environment, LockedSession) error) error { + s.record("WithEnrollment", s.enrollment != nil, environment, credentialHash) + return apply(ctx, s.enrollment, s.environment, s.locked) +} + +func (s *fakeDeviceStorage) WithArtifactStaging(context.Context, ArtifactStagingKey, func(context.Context, ArtifactStagingTx) error) error { + s.t.Fatal("unexpected call to WithArtifactStaging") + return nil +} + +func (s *fakeDeviceStorage) DeleteSessionArtifact(context.Context, string, string, string) error { + s.t.Fatal("unexpected call to DeleteSessionArtifact") + return nil +} + +func deviceService(t *testing.T, storage *fakeDeviceStorage) *Service { + t.Helper() + service, err := NewService(storage) + if err != nil { + t.Fatal(err) + } + return service +} + +const credentialDigest = "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08" + +func TestNewDeviceRegistration(t *testing.T) { + registration, err := NewDeviceRegistration(" runtime ", strings.ToUpper(credentialDigest)) + if err != nil || registration != (DeviceRegistration{Name: "runtime", CredentialHash: credentialDigest}) { + t.Fatalf("registration %+v, %v", registration, err) + } + if _, err := NewDeviceRegistration(strings.Repeat("n", 256), credentialDigest); err != nil { + t.Fatalf("longest name: %v", err) + } + for name, input := range map[string][2]string{ + "blank name": {" ", credentialDigest}, + "long name": {strings.Repeat("n", 257), credentialDigest}, + "not hex": {"runtime", "not-a-digest"}, + "short digest": {"runtime", credentialDigest[:62]}, + "raw credential": {"runtime", "secret"}, + "missing digest": {"runtime", ""}, + "padded digest": {"runtime", " " + credentialDigest}, + "too long digest": {"runtime", credentialDigest + "00"}, + } { + if _, err := NewDeviceRegistration(input[0], input[1]); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("%s: %v", name, err) + } + } +} + +func TestDeviceUseCases(t *testing.T) { + storage := &fakeDeviceStorage{t: t} + if _, err := deviceService(t, storage).CreateDevice(t.Context(), "tenant", "runtime", "secret"); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("invalid registration: %v", err) + } + created := ExecutionDevice{ID: "device", Name: "runtime"} + storage.createDevice = func(DeviceRegistration) (ExecutionDevice, error) { return created, nil } + storage.revokeDevice = done + if device, err := deviceService(t, storage).CreateDevice(t.Context(), "tenant", " runtime ", strings.ToUpper(credentialDigest)); err != nil || device != created { + t.Fatalf("created %+v, %v", device, err) + } + if err := deviceService(t, storage).RevokeDevice(t.Context(), "tenant", "device"); err != nil { + t.Fatal(err) + } + want := []string{"CreateDevice tenant runtime " + credentialDigest, "RevokeDevice tenant device"} + if strings.Join(storage.calls, "\n") != strings.Join(want, "\n") { + t.Fatalf("calls %q", storage.calls) + } + + // A heartbeat reports whether the device, or the credential it was + // authenticated with, still has authority. + for _, current := range []bool{true, false} { + storage := &fakeDeviceStorage{t: t, touchDevice: returns(current), touchAuthenticatedDevice: returns(current)} + service := deviceService(t, storage) + runtime, err := service.TouchRuntimeHeartbeat(t.Context(), "device") + if err != nil || runtime != (runtimedevice.HeartbeatStatus{Liveness: "online", Deleted: !current}) { + t.Fatalf("runtime heartbeat %+v, %v", runtime, err) + } + daemon, err := service.TouchAgentDaemonHeartbeat(t.Context(), runtimedevice.Heartbeat{RuntimeID: "device", CredentialHash: credentialDigest}) + if err != nil || daemon != runtime { + t.Fatalf("daemon heartbeat %+v, %v", daemon, err) + } + if want := []string{"TouchDevice device", "TouchAuthenticatedDevice device " + credentialDigest}; strings.Join(storage.calls, "\n") != strings.Join(want, "\n") { + t.Fatalf("calls %q", storage.calls) + } + } + failing := func() (bool, error) { return true, errStorage } + storage = &fakeDeviceStorage{t: t, touchDevice: failing, touchAuthenticatedDevice: failing} + if status, err := deviceService(t, storage).TouchRuntimeHeartbeat(t.Context(), "device"); !errors.Is(err, errStorage) || status != (runtimedevice.HeartbeatStatus{}) { + t.Fatalf("failed runtime heartbeat %+v, %v", status, err) + } + if status, err := deviceService(t, storage).TouchAgentDaemonHeartbeat(t.Context(), runtimedevice.Heartbeat{RuntimeID: "device"}); !errors.Is(err, errStorage) || status != (runtimedevice.HeartbeatStatus{}) { + t.Fatalf("failed daemon heartbeat %+v, %v", status, err) + } +} + +func TestEnrollRuntimeBindsUnderTheSessionLock(t *testing.T) { + environment := Environment{ID: "environment", SessionID: "session"} + authorized := returns(EnrollmentAuthority{KeyID: "key", WorkspaceDirectory: "/workspace"}) + for _, test := range []struct { + name string + tx fakeTx + locked LockedSession + want error + calls []string + }{ + {"enrolls", fakeTx{authorizeEnrollment: authorized, enrollDevice: returns("device"), bindDevice: done}, + LockedSession{}, nil, []string{"AuthorizeEnrollment", "EnrollDevice key", "BindDevice device"}}, + {"deleted Session", fakeTx{}, LockedSession{Deleted: true}, ErrNotFound, nil}, + {"revoked key", fakeTx{authorizeEnrollment: func() (EnrollmentAuthority, error) { return EnrollmentAuthority{}, ErrNotFound }}, + LockedSession{}, ErrNotFound, []string{"AuthorizeEnrollment"}}, + {"device of another key", fakeTx{authorizeEnrollment: authorized, enrollDevice: func() (string, error) { return "", ErrDeviceBindingConflict }}, + LockedSession{}, ErrDeviceBindingConflict, []string{"AuthorizeEnrollment", "EnrollDevice key"}}, + {"Session bound elsewhere", fakeTx{authorizeEnrollment: authorized, enrollDevice: returns("device"), bindDevice: func() error { return ErrDeviceBindingConflict }}, + LockedSession{}, ErrDeviceBindingConflict, []string{"AuthorizeEnrollment", "EnrollDevice key", "BindDevice device"}}, + } { + t.Run(test.name, func(t *testing.T) { + tx := test.tx + tx.t = t + storage := &fakeDeviceStorage{t: t, enrollment: &tx, environment: environment, locked: test.locked} + enrolled, err := deviceService(t, storage).EnrollRuntime(t.Context(), "environment", credentialDigest) + if test.want == nil && err != nil || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v, want %v", err, test.want) + } + if want := (RuntimeEnrollment{DeviceID: "device", SessionID: "session", EnvironmentID: "environment", WorkspaceDirectory: "/workspace"}); test.want == nil && enrolled != want { + t.Fatalf("enrollment %+v", enrolled) + } + if test.want != nil && enrolled != (RuntimeEnrollment{}) { + t.Fatalf("failed enrollment returned %+v", enrolled) + } + if storage.calls[0] != "WithEnrollment environment "+credentialDigest { + t.Fatalf("storage calls %q", storage.calls) + } + assertCalls(t, &tx, test.calls...) + }) + } +} diff --git a/services/core/internal/sessions/environment.go b/services/core/internal/sessions/environment.go index 59b1bc497..745ecc791 100644 --- a/services/core/internal/sessions/environment.go +++ b/services/core/internal/sessions/environment.go @@ -1,13 +1,66 @@ package sessions import ( + "context" "encoding/hex" "encoding/json" + "fmt" "time" "github.com/google/uuid" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" ) +// EnvironmentReader reads Environments, their preparation and the Session +// data frozen for it. +type EnvironmentReader interface { + // GetEnvironment reads the tenant's Environment of a Session that was not + // publicly deleted; otherwise it is ErrNotFound. + GetEnvironment(ctx context.Context, tenant, environment string) (Environment, error) + // GetSessionEnvironment reads the Environment of the tenant's Session + // that was not publicly deleted; otherwise it is ErrNotFound. + GetSessionEnvironment(ctx context.Context, tenant, session string) (Environment, error) + // ListEnvironmentInitializations lists, in Environment order after the + // given Environment, a page of the live Environments whose preparation is + // pending or running. + ListEnvironmentInitializations(ctx context.Context, after string) ([]EnvironmentInitialization, error) + // ReadEnvironmentSetup opens the setup frozen for the Session's + // Environment. A missing Session is ErrNotFound and a missing credential + // key credentialcrypto.ErrUnavailable; frozen data that does not open or + // validate is an internal error. + ReadEnvironmentSetup(ctx context.Context, tenant, session string) (environmentconfig.Setup, error) + // ReadInitialEnvironmentFile opens the initial file frozen at position for + // the Session's Environment, so an installation holds one file at a time. + // A missing file is ErrNotFound and a missing credential key + // credentialcrypto.ErrUnavailable; a file that does not open or match its + // recorded size is an internal error. + ReadInitialEnvironmentFile(ctx context.Context, tenant, session string, position int) (environmentconfig.InitialFileMetadata, []byte, error) +} + +// CreatesEnvironment reports whether a Session created with the configuration +// snapshot has an Environment: a self_hosted or openai_hosted one does, and +// none or no Environment does not. Any other snapshot is ErrInvalidInput. +func CreatesEnvironment(configuration json.RawMessage) (bool, error) { + var snapshot struct { + Environment *struct { + Type string `json:"type"` + } `json:"environment"` + } + if err := json.Unmarshal(configuration, &snapshot); err != nil { + return false, fmt.Errorf("%w: invalid environment configuration", ErrInvalidInput) + } + if snapshot.Environment == nil || snapshot.Environment.Type == "none" { + return false, nil + } + switch snapshot.Environment.Type { + case "self_hosted", "openai_hosted": + return true, nil + default: + return false, fmt.Errorf("%w: unsupported environment type", ErrInvalidInput) + } +} + // Environment retains execution ownership; its configuration is an internal snapshot, not a public response. type Environment struct { Initialization string diff --git a/services/core/internal/sessions/environment_test.go b/services/core/internal/sessions/environment_test.go new file mode 100644 index 000000000..a3df7db70 --- /dev/null +++ b/services/core/internal/sessions/environment_test.go @@ -0,0 +1,25 @@ +package sessions + +import ( + "errors" + "testing" +) + +func TestCreatesEnvironment(t *testing.T) { + for configuration, want := range map[string]bool{ + `{}`: false, + `{"environment":null}`: false, + `{"environment":{"type":"none"}}`: false, + `{"environment":{"type":"self_hosted"}}`: true, + `{"environment":{"type":"openai_hosted"}}`: true, + } { + if got, err := CreatesEnvironment([]byte(configuration)); err != nil || got != want { + t.Fatalf("%s: %v, %v", configuration, got, err) + } + } + for _, configuration := range []string{`not json`, `{"environment":{"type":"docker"}}`, `{"environment":{}}`} { + if _, err := CreatesEnvironment([]byte(configuration)); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("%s: %v", configuration, err) + } + } +} diff --git a/services/core/internal/sessions/execution.go b/services/core/internal/sessions/execution.go index 6432a8e5d..e222dc62c 100644 --- a/services/core/internal/sessions/execution.go +++ b/services/core/internal/sessions/execution.go @@ -18,6 +18,7 @@ func NewExecutionOperations(storage ExecutionStorage) (*ExecutionOperations, err // ExecutionStorage is the lease-bound Session storage, one family per line. type ExecutionStorage interface { + EnvironmentExecution FunctionExecution JournalExecution } diff --git a/services/core/internal/sessions/execution_environment.go b/services/core/internal/sessions/execution_environment.go new file mode 100644 index 000000000..764fa045a --- /dev/null +++ b/services/core/internal/sessions/execution_environment.go @@ -0,0 +1,355 @@ +package sessions + +import ( + "context" + "errors" + + "github.com/google/uuid" +) + +// EnvironmentKey names one Environment of a tenant. +type EnvironmentKey struct{ TenantID, EnvironmentID string } + +// EnvironmentConnection is the connection generation that observes an +// Environment and the last revision it committed. +type EnvironmentConnection struct { + Generation string + Revision int64 +} + +// EnvironmentExecution is the lease-bound storage of the Environment +// execution operations. +type EnvironmentExecution interface { + // WithInitialization runs apply in the transaction of the tenant's + // Session, with what the Session lock shows. A malformed Session ID is a + // missing Session. + WithInitialization(ctx context.Context, tenant, session string, apply func(context.Context, InitializationTx, LockedSession) error) error + // WithConnection runs apply in the transaction of the Session that owns + // the tenant's Environment, with what the Session lock shows. An + // Environment of a publicly deleted Session, or a missing one, is + // ErrNotFound. + WithConnection(ctx context.Context, tenant, environment string, apply func(context.Context, ConnectionTx, LockedSession) error) error + // WithDeviceBinding runs apply in the transaction of the tenant's Session, + // with what the Session lock shows. + WithDeviceBinding(ctx context.Context, tenant, session string, apply func(context.Context, DeviceBindingTx, LockedSession) error) error + // ListEnvironmentConnections lists, in Environment order after the given + // Environment, or from the first one when after is empty, a page of the + // Environments of Sessions that were not publicly deleted and that are + // connected or have a connection generation. + ListEnvironmentConnections(ctx context.Context, after string) ([]EnvironmentKey, error) +} + +// InitializationTx is the Session transaction the Environment initialization +// operations run in. +type InitializationTx interface { + EnvironmentFailureTx + // LoadEnvironment reads the Session's Environment. + LoadEnvironment(ctx context.Context) (Environment, error) + // LoadSessionDevice reads the Runtime device bound to the Session and + // reports whether the Session has one that still has authority. + LoadSessionDevice(ctx context.Context) (ExecutionDevice, bool, error) + // ClaimInitialization starts the pending preparation of the live + // Environment and reports whether it did. + ClaimInitialization(ctx context.Context, environment string) (bool, error) + // CompleteInitialization completes the running preparation of the live + // Environment and reports whether it did. + CompleteInitialization(ctx context.Context, environment string) (bool, error) + // FailInitialization marks the Environment's preparation failed unless it + // completed. + FailInitialization(ctx context.Context, environment string) error +} + +// ConnectionTx is the Session transaction the Environment connection +// operations run in. +type ConnectionTx interface { + InputActivityTx + // LoadEnvironment reads the Session's Environment. + LoadEnvironment(ctx context.Context) (Environment, error) + // LoadConnection reads the Environment's connection generation and + // reports whether it has one. + LoadConnection(ctx context.Context, environment string) (EnvironmentConnection, bool, error) + // ReplaceConnection starts the Environment's connection generation at + // revision zero, replacing any earlier one. + ReplaceConnection(ctx context.Context, environment, generation string) error + // AdvanceConnection records the revision the current generation + // committed. + AdvanceConnection(ctx context.Context, environment string, revision int64) error + // DeleteConnection discards the Environment's connection generation. + DeleteConnection(ctx context.Context, environment string) error + // SetConnectionStatus sets the Environment's connection status. + SetConnectionStatus(ctx context.Context, environment, status string) error +} + +// DeviceBindingTx is the Session transaction BindSessionDevice runs in. +type DeviceBindingTx interface { + // LoadDevice reports whether device is a device of the Session's tenant + // that was not revoked. + LoadDevice(ctx context.Context, device string) (bool, error) + // BindDevice binds the device to the Session. A Session bound to another + // device, or a device dedicated to another Session's Environment, is + // ErrDeviceBindingConflict. + BindDevice(ctx context.Context, device string) error +} + +// terminalEnvironment reports whether an Environment of status failed or +// expired, after which nothing prepares or connects it. +func terminalEnvironment(status string) bool { + return status == "failed" || status == "expired" +} + +// checkInitializationOwner confirms that the claim of a listed initialization +// still names the Session's live Environment: another Environment is +// ErrDeviceBindingConflict, and a failed or expired one ErrNotFound. +func checkInitializationOwner(owner EnvironmentInitialization, environment Environment) error { + if environment.ID != owner.EnvironmentID { + return ErrDeviceBindingConflict + } + if terminalEnvironment(environment.Status) { + return ErrNotFound + } + return nil +} + +// checkInitializationDevice confirms that the device the initialization was +// listed with is still the Session's device for this Environment. A Session +// without a device is ErrNotFound and another binding ErrTurnConflict. +func checkInitializationDevice(owner EnvironmentInitialization, environment Environment, bound ExecutionDevice, found bool) error { + if !found { + return ErrNotFound + } + if bound.ID != owner.DeviceID || bound.EnvironmentID != environment.ID { + return ErrTurnConflict + } + return nil +} + +// withInitialization runs apply on the Session's live Environment that owner +// names, in the transaction of its Session. +func (o *ExecutionOperations) withInitialization(ctx context.Context, owner EnvironmentInitialization, apply func(context.Context, InitializationTx, Environment) error) error { + return o.storage.WithInitialization(ctx, owner.TenantID, owner.SessionID, func(ctx context.Context, tx InitializationTx, locked LockedSession) error { + if err := locked.Public(); err != nil { + return err + } + environment, err := tx.LoadEnvironment(ctx) + if err != nil { + return err + } + if err := checkInitializationOwner(owner, environment); err != nil { + return err + } + return apply(ctx, tx, environment) + }) +} + +// advanceInitialization checks that the listed device still owns the +// preparation, then runs one preparation transition, which reports whether it +// applied; one that did not is ErrTurnConflict. +func (o *ExecutionOperations) advanceInitialization(ctx context.Context, owner EnvironmentInitialization, transition func(context.Context, InitializationTx, Environment) (bool, error)) error { + return o.withInitialization(ctx, owner, func(ctx context.Context, tx InitializationTx, environment Environment) error { + bound, found, err := tx.LoadSessionDevice(ctx) + if err != nil { + return err + } + if err := checkInitializationDevice(owner, environment, bound, found); err != nil { + return err + } + applied, err := transition(ctx, tx, environment) + if err == nil && !applied { + return ErrTurnConflict + } + return err + }) +} + +// ClaimEnvironmentInitialization starts the listed pending preparation on its +// device. +func (o *ExecutionOperations) ClaimEnvironmentInitialization(ctx context.Context, owner EnvironmentInitialization) error { + return o.advanceInitialization(ctx, owner, func(ctx context.Context, tx InitializationTx, environment Environment) (bool, error) { + return tx.ClaimInitialization(ctx, environment.ID) + }) +} + +// CompleteEnvironmentInitialization completes the claimed preparation. +func (o *ExecutionOperations) CompleteEnvironmentInitialization(ctx context.Context, owner EnvironmentInitialization) error { + return o.advanceInitialization(ctx, owner, func(ctx context.Context, tx InitializationTx, environment Environment) (bool, error) { + return tx.CompleteInitialization(ctx, environment.ID) + }) +} + +// FailEnvironmentInitialization records the failure of a preparation that did +// not complete and fails the Environment, as FailEnvironment reports it. It +// records state and settles work; it never destroys compute or a workspace. +func (o *ExecutionOperations) FailEnvironmentInitialization(ctx context.Context, owner EnvironmentInitialization, failure ProvisioningFailure) error { + return o.withInitialization(ctx, owner, func(ctx context.Context, tx InitializationTx, environment Environment) error { + if environment.Initialization == "complete" { + return ErrTurnConflict + } + if err := tx.FailInitialization(ctx, environment.ID); err != nil { + return err + } + return FailEnvironment(ctx, tx, environment, failure.Reason(), failure.Detail()) + }) +} + +// connectionGeneration returns the canonical form of a connection generation, +// a non-nil UUID, and reports whether generation is one. +func connectionGeneration(generation string) (string, bool) { + id, err := uuid.Parse(generation) + if err != nil || id == uuid.Nil { + return "", false + } + return id.String(), true +} + +// connectionStatus is the Environment status a connection observation +// reports. +func connectionStatus(connected bool) string { + if connected { + return "connected" + } + return "disconnected" +} + +// withConnection runs apply on the Environment in the transaction of its +// Session, and reports the input activity that apply changes. +func (o *ExecutionOperations) withConnection(ctx context.Context, tenant, environment string, apply func(context.Context, ConnectionTx, Environment) error) error { + return o.storage.WithConnection(ctx, tenant, environment, func(ctx context.Context, tx ConnectionTx, locked LockedSession) error { + if err := locked.Public(); err != nil { + return err + } + current, err := tx.LoadEnvironment(ctx) + if err != nil { + return err + } + return TrackInputActivity(ctx, tx, func(ctx context.Context) error { return apply(ctx, tx, current) }) + }) +} + +// recordConnection sets the Environment's connection status and journals the +// matching agent.session.environment. change. +func recordConnection(ctx context.Context, tx ConnectionTx, environment Environment, status string) error { + kind, err := EnvironmentType(environment.Configuration) + if err != nil { + return err + } + if err := tx.SetConnectionStatus(ctx, environment.ID, status); err != nil { + return err + } + return tx.AppendChanges(ctx, EnvironmentStateChange(environment.ID, kind, status)) +} + +// ReplaceEnvironmentConnection starts an ordered connection generation for +// the Environment without claiming a connection; a connected Environment +// becomes disconnected. The caller serializes replacements and never retries +// an older replacement after its successor. A failed or expired Environment is +// ErrInvalidInput. +func (o *ExecutionOperations) ReplaceEnvironmentConnection(ctx context.Context, tenant, environment, generation string) error { + generation, ok := connectionGeneration(generation) + if !ok { + return ErrInvalidInput + } + return o.withConnection(ctx, tenant, environment, func(ctx context.Context, tx ConnectionTx, current Environment) error { + if terminalEnvironment(current.Status) { + return ErrInvalidInput + } + previous, found, err := tx.LoadConnection(ctx, current.ID) + if err != nil { + return err + } + if found && previous.Generation == generation { + return nil + } + if err := tx.ReplaceConnection(ctx, current.ID, generation); err != nil { + return err + } + if current.Status == "connected" { + return recordConnection(ctx, tx, current, "disconnected") + } + return nil + }) +} + +// ObserveEnvironmentConnection commits only a newer observation of the +// current generation, with the status change it reports. Connectivity is +// transport evidence, not native preparation readiness or process quiescence. +// An observation of a failed or expired Environment is ErrInvalidInput. +func (o *ExecutionOperations) ObserveEnvironmentConnection(ctx context.Context, tenant, environment, generation string, revision int64, connected bool) error { + generation, ok := connectionGeneration(generation) + if !ok || revision <= 0 { + return ErrInvalidInput + } + return o.withConnection(ctx, tenant, environment, func(ctx context.Context, tx ConnectionTx, current Environment) error { + connection, found, err := tx.LoadConnection(ctx, current.ID) + if err != nil || !found { + return err + } + if connection.Generation != generation || connection.Revision >= revision { + return nil + } + if terminalEnvironment(current.Status) { + return ErrInvalidInput + } + if err := tx.AdvanceConnection(ctx, current.ID, revision); err != nil { + return err + } + status := connectionStatus(connected) + if current.Status == status { + return nil + } + return recordConnection(ctx, tx, current, status) + }) +} + +// ReconcileEnvironmentConnections runs before a new execution owner's +// connection producers start. It discards the generations of earlier +// processes and records that their connected transport was lost. An +// Environment that is gone by the time its turn comes is skipped. +func (o *ExecutionOperations) ReconcileEnvironmentConnections(ctx context.Context) error { + after := "" + for { + keys, err := o.storage.ListEnvironmentConnections(ctx, after) + if err != nil { + return err + } + if len(keys) == 0 { + return nil + } + for _, key := range keys { + err := o.withConnection(ctx, key.TenantID, key.EnvironmentID, func(ctx context.Context, tx ConnectionTx, current Environment) error { + if err := tx.DeleteConnection(ctx, current.ID); err != nil { + return err + } + if current.Status == "connected" { + return recordConnection(ctx, tx, current, "disconnected") + } + return nil + }) + if err != nil && !errors.Is(err, ErrNotFound) { + return err + } + after = key.EnvironmentID + } + } +} + +// BindSessionDevice binds the tenant's device to the Session, including a +// publicly deleted one. A retry with the same device succeeds and another +// device is ErrDeviceBindingConflict, so a Session's filesystem never moves +// silently. A malformed device ID is ErrInvalidInput, and a revoked or unknown +// device ErrNotFound. Dispatchers read the full Session binding before +// delivery. +func (o *ExecutionOperations) BindSessionDevice(ctx context.Context, tenant, session, device string) error { + if !validID(device) { + return ErrInvalidInput + } + return o.storage.WithDeviceBinding(ctx, tenant, session, func(ctx context.Context, tx DeviceBindingTx, _ LockedSession) error { + found, err := tx.LoadDevice(ctx, device) + if err != nil { + return err + } + if !found { + return ErrNotFound + } + return tx.BindDevice(ctx, device) + }) +} diff --git a/services/core/internal/sessions/execution_environment_test.go b/services/core/internal/sessions/execution_environment_test.go new file mode 100644 index 000000000..73461e26b --- /dev/null +++ b/services/core/internal/sessions/execution_environment_test.go @@ -0,0 +1,380 @@ +package sessions + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + "time" +) + +func (f *fakeTx) LoadSessionDevice(context.Context) (ExecutionDevice, bool, error) { + f.record("LoadSessionDevice", f.loadSessionDevice != nil) + return f.loadSessionDevice() +} + +func (f *fakeTx) ClaimInitialization(_ context.Context, environment string) (bool, error) { + f.record("ClaimInitialization", f.claimInitialization != nil, environment) + return f.claimInitialization() +} + +func (f *fakeTx) CompleteInitialization(_ context.Context, environment string) (bool, error) { + f.record("CompleteInitialization", f.completeInitialization != nil, environment) + return f.completeInitialization() +} + +func (f *fakeTx) FailInitialization(_ context.Context, environment string) error { + f.record("FailInitialization", f.failInitialization != nil, environment) + return f.failInitialization() +} + +func (f *fakeTx) LoadConnection(_ context.Context, environment string) (EnvironmentConnection, bool, error) { + f.record("LoadConnection", f.loadConnection != nil, environment) + return f.loadConnection() +} + +func (f *fakeTx) ReplaceConnection(_ context.Context, environment, generation string) error { + f.record("ReplaceConnection", f.replaceConnection != nil, environment, generation) + return f.replaceConnection() +} + +func (f *fakeTx) AdvanceConnection(_ context.Context, environment string, revision int64) error { + f.record("AdvanceConnection", f.advanceConnection != nil, environment, fmt.Sprint(revision)) + return f.advanceConnection() +} + +func (f *fakeTx) DeleteConnection(_ context.Context, environment string) error { + f.record("DeleteConnection", f.deleteConnection != nil, environment) + return f.deleteConnection() +} + +func (f *fakeTx) SetConnectionStatus(_ context.Context, environment, status string) error { + f.record("SetConnectionStatus", f.setConnectionStatus != nil, environment, status) + return f.setConnectionStatus() +} + +func (f *fakeTx) LoadDevice(_ context.Context, device string) (bool, error) { + f.record("LoadDevice", f.loadDevice != nil, device) + return f.loadDevice() +} + +func (f *fakeTx) BindDevice(_ context.Context, device string) error { + f.record("BindDevice", f.bindDevice != nil, device) + return f.bindDevice() +} + +func (f *fakeTx) AuthorizeEnrollment(context.Context) (EnrollmentAuthority, error) { + f.record("AuthorizeEnrollment", f.authorizeEnrollment != nil) + return f.authorizeEnrollment() +} + +func (f *fakeTx) EnrollDevice(_ context.Context, key string) (string, error) { + f.record("EnrollDevice", f.enrollDevice != nil, key) + return f.enrollDevice() +} + +// environmentTx is the transaction the Environment family's With methods +// apply in; a call without one fails the test. +func (s *fakeExecutionStorage) environmentTx(name string) *fakeTx { + s.t.Helper() + if s.tx == nil { + s.t.Fatalf("unexpected call to %s", name) + } + return s.tx +} + +func (s *fakeExecutionStorage) WithInitialization(ctx context.Context, tenant, session string, apply func(context.Context, InitializationTx, LockedSession) error) error { + s.environmentTx("WithInitialization").record("WithInitialization", true, tenant, session) + return apply(ctx, s.tx, s.locked) +} + +func (s *fakeExecutionStorage) WithConnection(ctx context.Context, tenant, environment string, apply func(context.Context, ConnectionTx, LockedSession) error) error { + s.environmentTx("WithConnection").record("WithConnection", true, tenant, environment) + if err := s.connections[environment]; err != nil { + return err + } + return apply(ctx, s.tx, s.locked) +} + +func (s *fakeExecutionStorage) WithDeviceBinding(ctx context.Context, tenant, session string, apply func(context.Context, DeviceBindingTx, LockedSession) error) error { + s.environmentTx("WithDeviceBinding").record("WithDeviceBinding", true, tenant, session) + return apply(ctx, s.tx, s.locked) +} + +func (s *fakeExecutionStorage) ListEnvironmentConnections(_ context.Context, after string) ([]EnvironmentKey, error) { + s.environmentTx("ListEnvironmentConnections").record("ListEnvironmentConnections", s.pages != nil, after) + var page []EnvironmentKey + if len(s.pages) > 0 { + page, s.pages = s.pages[0], s.pages[1:] + } + return page, nil +} + +// loads is a fake method that reads value and reports whether it was found. +func loads[T any](value T, found bool) func() (T, bool, error) { + return func() (T, bool, error) { return value, found, nil } +} + +func executionOperations(t *testing.T, storage *fakeExecutionStorage) *ExecutionOperations { + t.Helper() + storage.t = t + operations, err := NewExecutionOperations(storage) + if err != nil { + t.Fatal(err) + } + return operations +} + +var hostedEnvironment = Environment{ID: "environment", SessionID: "session", Status: "pending", Configuration: []byte(`{"type":"openai_hosted"}`)} + +func TestInitializationOwnerAndDeviceRules(t *testing.T) { + owner := EnvironmentInitialization{EnvironmentID: "environment", DeviceID: "device"} + for _, test := range []struct { + name string + environment Environment + want error + }{ + {"live", Environment{ID: "environment", Status: "pending"}, nil}, + {"another Environment", Environment{ID: "other", Status: "pending"}, ErrDeviceBindingConflict}, + {"failed", Environment{ID: "environment", Status: "failed"}, ErrNotFound}, + {"expired", Environment{ID: "environment", Status: "expired"}, ErrNotFound}, + } { + if err := checkInitializationOwner(owner, test.environment); !errors.Is(err, test.want) { + t.Fatalf("owner %s: %v", test.name, err) + } + } + environment := Environment{ID: "environment"} + for _, test := range []struct { + name string + bound ExecutionDevice + found bool + want error + }{ + {"listed device", ExecutionDevice{ID: "device", EnvironmentID: "environment"}, true, nil}, + {"no device", ExecutionDevice{}, false, ErrNotFound}, + {"another device", ExecutionDevice{ID: "other", EnvironmentID: "environment"}, true, ErrTurnConflict}, + {"another Environment's device", ExecutionDevice{ID: "device", EnvironmentID: "other"}, true, ErrTurnConflict}, + } { + if err := checkInitializationDevice(owner, environment, test.bound, test.found); !errors.Is(err, test.want) { + t.Fatalf("device %s: %v", test.name, err) + } + } +} + +func TestInitializationTransitionsRunUnderTheListedDevice(t *testing.T) { + owner := EnvironmentInitialization{TenantID: "tenant", SessionID: "session", EnvironmentID: "environment", DeviceID: "device"} + listed := func() (ExecutionDevice, bool, error) { + return ExecutionDevice{ID: "device", EnvironmentID: "environment"}, true, nil + } + checked := []string{"WithInitialization tenant session", "LoadEnvironment", "LoadSessionDevice"} + for _, test := range []struct { + name string + tx fakeTx + locked LockedSession + operate func(*ExecutionOperations) error + want error + calls []string + }{ + {"claim", fakeTx{loadEnvironment: returns(hostedEnvironment), loadSessionDevice: listed, claimInitialization: returns(true)}, + LockedSession{}, func(o *ExecutionOperations) error { return o.ClaimEnvironmentInitialization(t.Context(), owner) }, + nil, append(checked, "ClaimInitialization environment")}, + {"claim that did not apply", fakeTx{loadEnvironment: returns(hostedEnvironment), loadSessionDevice: listed, claimInitialization: returns(false)}, + LockedSession{}, func(o *ExecutionOperations) error { return o.ClaimEnvironmentInitialization(t.Context(), owner) }, + ErrTurnConflict, append(checked, "ClaimInitialization environment")}, + {"complete", fakeTx{loadEnvironment: returns(hostedEnvironment), loadSessionDevice: listed, completeInitialization: returns(true)}, + LockedSession{}, func(o *ExecutionOperations) error { return o.CompleteEnvironmentInitialization(t.Context(), owner) }, + nil, append(checked, "CompleteInitialization environment")}, + {"device moved", fakeTx{loadEnvironment: returns(hostedEnvironment), loadSessionDevice: loads(ExecutionDevice{ID: "other", EnvironmentID: "environment"}, true)}, + LockedSession{}, func(o *ExecutionOperations) error { return o.CompleteEnvironmentInitialization(t.Context(), owner) }, + ErrTurnConflict, checked}, + {"deleted Session", fakeTx{}, + LockedSession{Deleted: true}, func(o *ExecutionOperations) error { return o.ClaimEnvironmentInitialization(t.Context(), owner) }, + ErrNotFound, checked[:1]}, + {"failed Environment", fakeTx{loadEnvironment: returns(Environment{ID: "environment", Status: "failed"})}, + LockedSession{}, func(o *ExecutionOperations) error { return o.ClaimEnvironmentInitialization(t.Context(), owner) }, + ErrNotFound, checked[:2]}, + } { + t.Run(test.name, func(t *testing.T) { + tx := test.tx + tx.t = t + err := test.operate(executionOperations(t, &fakeExecutionStorage{tx: &tx, locked: test.locked})) + if test.want == nil && err != nil || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v, want %v", err, test.want) + } + assertCalls(t, &tx, test.calls...) + }) + } +} + +func TestFailEnvironmentInitializationKeepsACompletedPreparation(t *testing.T) { + owner := EnvironmentInitialization{TenantID: "tenant", SessionID: "session", EnvironmentID: "environment"} + failure := ProvisioningFailure{Step: ProvisioningHarness} + + complete := hostedEnvironment + complete.Initialization = "complete" + tx := fakeTx{t: t, loadEnvironment: returns(complete)} + err := executionOperations(t, &fakeExecutionStorage{tx: &tx}).FailEnvironmentInitialization(t.Context(), owner, failure) + if !errors.Is(err, ErrTurnConflict) { + t.Fatalf("completed preparation failed: %v", err) + } + assertCalls(t, &tx, "WithInitialization tenant session", "LoadEnvironment") + + // The preparation fails before the Environment does; the device is not + // consulted. + tx = fakeTx{t: t, loadEnvironment: returns(hostedEnvironment), failInitialization: done, + recordEnvironmentFailure: func() (time.Time, error) { return time.Time{}, errStorage }} + err = executionOperations(t, &fakeExecutionStorage{tx: &tx}).FailEnvironmentInitialization(t.Context(), owner, failure) + if !errors.Is(err, errStorage) { + t.Fatal(err) + } + assertCalls(t, &tx, "WithInitialization tenant session", "LoadEnvironment", "FailInitialization environment", + "RecordEnvironmentFailure environment "+failure.Reason()) +} + +// quietInput is a fake LoadEnvironmentInput with no input activity. +func quietInput() (*EnvironmentInputState, error) { return nil, nil } + +func TestConnectionGenerationsAreOrdered(t *testing.T) { + const generation = "0b7c6f0e-3d7b-4bd2-9f55-6f1d3f0f8a11" + connected := hostedEnvironment + connected.Status = "connected" + var changes []SessionChange + tracked := func(calls ...string) []string { + return append(append([]string{"WithConnection tenant environment", "LoadEnvironment", "LoadEnvironmentInput"}, calls...), "LoadEnvironmentInput") + } + for _, test := range []struct { + name string + tx fakeTx + locked LockedSession + operate func(*ExecutionOperations) error + want error + calls []string + }{ + {"replacement disconnects", fakeTx{loadEnvironment: returns(connected), loadEnvironmentInput: quietInput, loadConnection: loads(EnvironmentConnection{}, false), replaceConnection: done, setConnectionStatus: done, appendChanges: collect(&changes)}, LockedSession{}, + func(o *ExecutionOperations) error { + return o.ReplaceEnvironmentConnection(t.Context(), "tenant", "environment", strings.ToUpper(generation)) + }, + nil, tracked("LoadConnection environment", "ReplaceConnection environment "+generation, "SetConnectionStatus environment disconnected", "AppendChanges agent.session.environment.disconnected")}, + {"repeated replacement", fakeTx{loadEnvironment: returns(connected), loadEnvironmentInput: quietInput, loadConnection: loads(EnvironmentConnection{Generation: generation}, true)}, LockedSession{}, + func(o *ExecutionOperations) error { + return o.ReplaceEnvironmentConnection(t.Context(), "tenant", "environment", generation) + }, + nil, tracked("LoadConnection environment")}, + {"replacement of a failed Environment", fakeTx{loadEnvironment: returns(Environment{ID: "environment", Status: "failed"}), loadEnvironmentInput: quietInput}, LockedSession{}, + func(o *ExecutionOperations) error { + return o.ReplaceEnvironmentConnection(t.Context(), "tenant", "environment", generation) + }, + ErrInvalidInput, tracked()[:3]}, + {"newer observation connects", fakeTx{loadEnvironment: returns(hostedEnvironment), loadEnvironmentInput: quietInput, loadConnection: loads(EnvironmentConnection{Generation: generation, Revision: 1}, true), advanceConnection: done, setConnectionStatus: done, appendChanges: collect(&changes)}, LockedSession{}, + func(o *ExecutionOperations) error { + return o.ObserveEnvironmentConnection(t.Context(), "tenant", "environment", generation, 2, true) + }, + nil, tracked("LoadConnection environment", "AdvanceConnection environment 2", "SetConnectionStatus environment connected", "AppendChanges agent.session.environment.connected")}, + {"observation without a status change", fakeTx{loadEnvironment: returns(connected), loadEnvironmentInput: quietInput, loadConnection: loads(EnvironmentConnection{Generation: generation, Revision: 1}, true), advanceConnection: done}, LockedSession{}, + func(o *ExecutionOperations) error { + return o.ObserveEnvironmentConnection(t.Context(), "tenant", "environment", generation, 2, true) + }, + nil, tracked("LoadConnection environment", "AdvanceConnection environment 2")}, + {"stale observation", fakeTx{loadEnvironment: returns(hostedEnvironment), loadEnvironmentInput: quietInput, loadConnection: loads(EnvironmentConnection{Generation: generation, Revision: 2}, true)}, LockedSession{}, + func(o *ExecutionOperations) error { + return o.ObserveEnvironmentConnection(t.Context(), "tenant", "environment", generation, 2, true) + }, + nil, tracked("LoadConnection environment")}, + {"observation of a replaced generation", fakeTx{loadEnvironment: returns(hostedEnvironment), loadEnvironmentInput: quietInput, loadConnection: loads(EnvironmentConnection{Generation: "d2b0c1a4-4c83-4b7f-8f0a-3c9b54bb1d52", Revision: 1}, true)}, LockedSession{}, + func(o *ExecutionOperations) error { + return o.ObserveEnvironmentConnection(t.Context(), "tenant", "environment", generation, 5, false) + }, + nil, tracked("LoadConnection environment")}, + {"observation of a failed Environment", fakeTx{loadEnvironment: returns(Environment{ID: "environment", Status: "failed"}), loadEnvironmentInput: quietInput, loadConnection: loads(EnvironmentConnection{Generation: generation, Revision: 1}, true)}, LockedSession{}, + func(o *ExecutionOperations) error { + return o.ObserveEnvironmentConnection(t.Context(), "tenant", "environment", generation, 2, true) + }, + ErrInvalidInput, tracked("LoadConnection environment")[:4]}, + {"deleted Session", fakeTx{}, LockedSession{Deleted: true}, + func(o *ExecutionOperations) error { + return o.ObserveEnvironmentConnection(t.Context(), "tenant", "environment", generation, 2, true) + }, + ErrNotFound, []string{"WithConnection tenant environment"}}, + } { + t.Run(test.name, func(t *testing.T) { + tx := test.tx + tx.t = t + changes = nil + err := test.operate(executionOperations(t, &fakeExecutionStorage{tx: &tx, locked: test.locked})) + if test.want == nil && err != nil || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v, want %v", err, test.want) + } + assertCalls(t, &tx, test.calls...) + for _, change := range changes { + if state := change.Event.Environment; state == nil || state.ID != "environment" || state.Type != "openai_hosted" { + t.Fatalf("change %+v", change.Event) + } + } + }) + } + for _, invalid := range []struct { + generation string + revision int64 + }{{"not-a-uuid", 1}, {"00000000-0000-0000-0000-000000000000", 1}, {generation, 0}} { + tx := fakeTx{t: t} + operations := executionOperations(t, &fakeExecutionStorage{tx: &tx}) + if err := operations.ObserveEnvironmentConnection(t.Context(), "tenant", "environment", invalid.generation, invalid.revision, true); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("observation %+v: %v", invalid, err) + } + if invalid.revision > 0 { + if err := operations.ReplaceEnvironmentConnection(t.Context(), "tenant", "environment", invalid.generation); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("replacement %+v: %v", invalid, err) + } + } + assertCalls(t, &tx) + } +} + +func TestReconcileEnvironmentConnectionsDisconnectsEveryPage(t *testing.T) { + connected := hostedEnvironment + connected.Status = "connected" + var changes []SessionChange + tx := fakeTx{t: t, loadEnvironment: returns(connected), loadEnvironmentInput: quietInput, deleteConnection: done, setConnectionStatus: done, appendChanges: collect(&changes)} + storage := &fakeExecutionStorage{tx: &tx, pages: [][]EnvironmentKey{{{"tenant", "environment"}, {"tenant", "gone"}}, {{"tenant", "environment"}}}, connections: map[string]error{"gone": ErrNotFound}} + if err := executionOperations(t, storage).ReconcileEnvironmentConnections(t.Context()); err != nil { + t.Fatal(err) + } + reconcile := []string{"WithConnection tenant environment", "LoadEnvironment", "LoadEnvironmentInput", "DeleteConnection environment", + "SetConnectionStatus environment disconnected", "AppendChanges agent.session.environment.disconnected", "LoadEnvironmentInput"} + want := append([]string{"ListEnvironmentConnections "}, reconcile...) + want = append(want, "WithConnection tenant gone", "ListEnvironmentConnections gone") + want = append(want, reconcile...) + want = append(want, "ListEnvironmentConnections environment") + assertCalls(t, &tx, want...) + + tx = fakeTx{t: t} + storage = &fakeExecutionStorage{tx: &tx, pages: [][]EnvironmentKey{{{"tenant", "environment"}}}, connections: map[string]error{"environment": errStorage}} + if err := executionOperations(t, storage).ReconcileEnvironmentConnections(t.Context()); !errors.Is(err, errStorage) { + t.Fatalf("storage failure: %v", err) + } +} + +func TestBindSessionDeviceRequiresTheTenantsDevice(t *testing.T) { + // A malformed device is ErrInvalidInput before any storage call, whether or + // not the Session exists. + if err := executionOperations(t, &fakeExecutionStorage{}).BindSessionDevice(t.Context(), "tenant", "missing", "device"); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("malformed device: %v", err) + } + + device := "5b0cf2a4-6c41-4f55-9d2f-3f6f1b1c7e10" + tx := fakeTx{t: t, loadDevice: returns(false)} + operations := executionOperations(t, &fakeExecutionStorage{tx: &tx, locked: LockedSession{Deleted: true}}) + if err := operations.BindSessionDevice(t.Context(), "tenant", "session", device); !errors.Is(err, ErrNotFound) { + t.Fatalf("unknown device: %v", err) + } + assertCalls(t, &tx, "WithDeviceBinding tenant session", "LoadDevice "+device) + + // A publicly deleted Session still binds. + tx = fakeTx{t: t, loadDevice: returns(true), bindDevice: func() error { return ErrDeviceBindingConflict }} + operations = executionOperations(t, &fakeExecutionStorage{tx: &tx, locked: LockedSession{Deleted: true}}) + if err := operations.BindSessionDevice(t.Context(), "tenant", "session", device); !errors.Is(err, ErrDeviceBindingConflict) { + t.Fatalf("conflict: %v", err) + } + assertCalls(t, &tx, "WithDeviceBinding tenant session", "LoadDevice "+device, "BindDevice "+device) +} diff --git a/services/core/internal/sessions/execution_functions_test.go b/services/core/internal/sessions/execution_functions_test.go index a9b976046..f3c016138 100644 --- a/services/core/internal/sessions/execution_functions_test.go +++ b/services/core/internal/sessions/execution_functions_test.go @@ -11,11 +11,21 @@ import ( ) // fakeExecutionStorage is strict lease-bound Session storage: a method whose -// func is unset fails the test. +// func is unset fails the test. The Environment family's With methods record +// their call on tx, then run apply with tx and locked; without tx they fail +// the test. type fakeExecutionStorage struct { t *testing.T withFunctionTurn func(ctx context.Context, tenant, session, turn string, apply func(context.Context, FunctionTx, Turn) error) error withTurnJournal func(ctx context.Context, tenant, session string, apply func(context.Context, TurnJournalTx) error) error + + tx *fakeTx + locked LockedSession + // pages are the pages ListEnvironmentConnections reads, one per call. + pages [][]EnvironmentKey + // connections is the error each WithConnection call returns in place of + // applying, by Environment. + connections map[string]error } var _ ExecutionStorage = (*fakeExecutionStorage)(nil) diff --git a/services/core/internal/sessions/executor_credentials.go b/services/core/internal/sessions/executor_credentials.go index 311124c0c..c88505295 100644 --- a/services/core/internal/sessions/executor_credentials.go +++ b/services/core/internal/sessions/executor_credentials.go @@ -1,11 +1,21 @@ package sessions import ( + "context" "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" ) +// ExecutorCredentialReader authenticates Environment executor credentials. +type ExecutorCredentialReader interface { + // AuthenticateEnvironmentExecutor checks, in one database snapshot, that + // the credential whose SHA-256 digest is given is current for the + // Environment and belongs to its Session's creator, and returns the + // Session's tenant. Any other credential is ErrNotFound. + AuthenticateEnvironmentExecutor(ctx context.Context, environment, digest string) (string, error) +} + type IssuedExecutorCredential struct { KeyID string `json:"key_id"` EnvironmentID string `json:"environment_id,omitempty"` diff --git a/services/core/internal/sessions/reader.go b/services/core/internal/sessions/reader.go index 17b8357c0..5fe5d1e34 100644 --- a/services/core/internal/sessions/reader.go +++ b/services/core/internal/sessions/reader.go @@ -1,9 +1,12 @@ package sessions -// Reader reads a Session's public resources, one family per line. Callers use -// it directly; no use case forwards a read. +// Reader reads a Session's resources, Environment and devices, one family per +// line. Callers use it directly; no use case forwards a read. type Reader interface { ArtifactReader + DeviceReader + EnvironmentReader + ExecutorCredentialReader ItemReader SubagentReader } diff --git a/services/core/internal/sessions/service.go b/services/core/internal/sessions/service.go index 823a52804..d75300cbb 100644 --- a/services/core/internal/sessions/service.go +++ b/services/core/internal/sessions/service.go @@ -17,4 +17,5 @@ func NewService(storage Storage) (*Service, error) { // Storage persists the pooled Session use cases, one family per line. type Storage interface { ArtifactStorage + DeviceStorage } diff --git a/services/core/internal/sessions/transaction_test.go b/services/core/internal/sessions/transaction_test.go index 0a365afce..3cb35c1e8 100644 --- a/services/core/internal/sessions/transaction_test.go +++ b/services/core/internal/sessions/transaction_test.go @@ -38,6 +38,19 @@ type fakeTx struct { loadPendingFileWrite func() (bool, error) loadBoundDevice func() (bool, error) insertEnvironmentDevice func() error + loadSessionDevice func() (ExecutionDevice, bool, error) + claimInitialization func() (bool, error) + completeInitialization func() (bool, error) + failInitialization func() error + loadConnection func() (EnvironmentConnection, bool, error) + replaceConnection func() error + advanceConnection func() error + deleteConnection func() error + setConnectionStatus func() error + loadDevice func() (bool, error) + bindDevice func() error + authorizeEnrollment func() (EnrollmentAuthority, error) + enrollDevice func() (string, error) loadJournalTurn func() (JournalTurn, bool, error) matchEvents func() (bool, error) @@ -71,6 +84,10 @@ var ( _ TurnJournalTx = (*fakeTx)(nil) _ TurnEventTx = (*fakeTx)(nil) _ InputProjectionTx = (*fakeTx)(nil) + _ InitializationTx = (*fakeTx)(nil) + _ ConnectionTx = (*fakeTx)(nil) + _ DeviceBindingTx = (*fakeTx)(nil) + _ EnrollmentTx = (*fakeTx)(nil) ) func (f *fakeTx) record(name string, set bool, detail ...string) { diff --git a/services/core/internal/store/admin_session_archive_test.go b/services/core/internal/store/admin_session_archive_test.go index b8da724b5..35c6d4a1e 100644 --- a/services/core/internal/store/admin_session_archive_test.go +++ b/services/core/internal/store/admin_session_archive_test.go @@ -154,7 +154,7 @@ func TestManagedSessionArchiveRetainsHistoryAndSettledResources(t *testing.T) { if !reflect.DeepEqual(history, adminMutationSnapshot(t, s, "sessions", "turns", "session_items", "session_artifacts", "source_files", "pg_largeobject", "pg_largeobject_metadata")) { t.Fatal("archive changed persisted history or artifacts") } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + 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) { @@ -235,7 +235,7 @@ func TestManagedSessionArchivePreservesFailuresAndRejectsSelfHosted(t *testing.T if _, err := s.pool.Exec(t.Context(), "UPDATE environments SET initialization='running' WHERE id=$1", owner.EnvironmentID); err != nil { t.Fatal(err) } - if err := w.FailEnvironmentInitialization(t.Context(), sessions.EnvironmentInitialization{EnvironmentID: owner.EnvironmentID, SessionID: owner.SessionID, TenantID: owner.TenantID, DeviceID: owner.DeviceID}, sessions.ProvisioningFailure{Step: sessions.ProvisioningSetupCommand, Index: 0, ExitCode: 2}); err != nil { + if err := sessionExecution(t, w.lease).FailEnvironmentInitialization(t.Context(), sessions.EnvironmentInitialization{EnvironmentID: owner.EnvironmentID, SessionID: owner.SessionID, TenantID: owner.TenantID, DeviceID: owner.DeviceID}, sessions.ProvisioningFailure{Step: sessions.ProvisioningSetupCommand, Index: 0, ExitCode: 2}); err != nil { t.Fatal(err) } failed, err := s.GetSession(t.Context(), tenant, session.ID) diff --git a/services/core/internal/store/archive_cancellation_test.go b/services/core/internal/store/archive_cancellation_test.go index b891ebd2f..62280521a 100644 --- a/services/core/internal/store/archive_cancellation_test.go +++ b/services/core/internal/store/archive_cancellation_test.go @@ -74,19 +74,19 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { t.Fatal(err) } } - if err := writer.BindSessionDevice(t.Context(), project.TenantID, session.ID, owner.DeviceID); err != nil { + if err := leased.Sessions.BindSessionDevice(t.Context(), project.TenantID, session.ID, owner.DeviceID); err != nil { t.Fatal(err) } generation := uuid.NewString() - if err := writer.ReplaceEnvironmentConnection(t.Context(), project.TenantID, session.Environment.ID, generation); err != nil { + if err := leased.Sessions.ReplaceEnvironmentConnection(t.Context(), project.TenantID, session.Environment.ID, generation); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), project.TenantID, session.Environment.ID, generation, 1, true); err != nil { + if err := leased.Sessions.ObserveEnvironmentConnection(t.Context(), project.TenantID, session.Environment.ID, generation, 1, true); err != nil { t.Fatal(err) } server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - handler, registry, err := runtime.NewGateway(s, wsURL) + handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), s, wsURL) if err != nil { t.Fatal(err) } @@ -147,14 +147,14 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { } } if scenario == "revoke_before_archive" || scenario == "cancel_revoke_archive" { - if err := s.RevokeDevice(t.Context(), h.tenant, owner.DeviceID); err != nil { + if err := fixtureSessionService(t, db).RevokeDevice(t.Context(), h.tenant, owner.DeviceID); err != nil { t.Fatal(err) } } var revokeDone chan error if scenario == "revoke_concurrent_archive" { revokeDone = make(chan error, 1) - go func() { revokeDone <- s.RevokeDevice(t.Context(), h.tenant, owner.DeviceID) }() + go func() { revokeDone <- fixtureSessionService(t, db).RevokeDevice(t.Context(), h.tenant, owner.DeviceID) }() } archived, err := writer.ArchiveManagedSession(auditCtx, h.tenant, session.ID, 1) @@ -167,7 +167,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { } } if scenario == "revoke_after_archive" { - if err := s.RevokeDevice(t.Context(), h.tenant, owner.DeviceID); err != nil { + if err := fixtureSessionService(t, db).RevokeDevice(t.Context(), h.tenant, owner.DeviceID); err != nil { t.Fatal(err) } } @@ -185,7 +185,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { if err != nil || current.Status != sessions.TurnWaiting || current.CancelRequestedAt.IsZero() { t.Fatal("archive must request rather than invent cancellation", current, err) } - if _, err := runtimegateway.NewAuthenticator(s).AuthenticateBearer(t.Context(), owner.DeviceID, secret); !errors.Is(err, runtimegateway.ErrAuthUnknownDevice) { + if _, err := runtimegateway.NewAuthenticator(fixtureSessionStore(db)).AuthenticateBearer(t.Context(), owner.DeviceID, secret); !errors.Is(err, runtimegateway.ErrAuthUnknownDevice) { t.Fatal("archive allowed renewed authority", err) } rejected, response, dialErr := websocket.DefaultDialer.Dial(u.String(), http.Header{"Authorization": {"Bearer " + secret}}) diff --git a/services/core/internal/store/claude_execution_test.go b/services/core/internal/store/claude_execution_test.go index 2d5503acb..3a94c3829 100644 --- a/services/core/internal/store/claude_execution_test.go +++ b/services/core/internal/store/claude_execution_test.go @@ -23,7 +23,7 @@ func claudeSession(t *testing.T, h *dispatchHarness, configuration string, prebo t.Fatal(err) } if prebound { - if err := h.s.BindSessionDevice(t.Context(), h.tenant, h.session.ID, h.device.ID); err != nil { + if err := bindSessionDevice(t, h.db, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } } @@ -73,7 +73,7 @@ func TestClaudeWorkerSelectsStoredEngineAndRestrictiveCapabilities(t *testing.T) t.Fatal(turn, err) } if !prebound { - if _, err := h.s.GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := fixtureSessionStore(h.db).GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("bound an incapable device", err) } } diff --git a/services/core/internal/store/claude_mcp_test.go b/services/core/internal/store/claude_mcp_test.go index ccab5314a..b72536d98 100644 --- a/services/core/internal/store/claude_mcp_test.go +++ b/services/core/internal/store/claude_mcp_test.go @@ -60,7 +60,7 @@ func TestClaudeMCPWaitsForCapableRuntime(t *testing.T) { t.Fatal("incapable runtime claimed work", turn, err) } if !prebound { - if _, err := h.s.GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := fixtureSessionStore(h.db).GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("bound an incapable runtime", err) } } diff --git a/services/core/internal/store/creation_stream_settlement_public_test.go b/services/core/internal/store/creation_stream_settlement_public_test.go index 1c4c23c19..e50fd5360 100644 --- a/services/core/internal/store/creation_stream_settlement_public_test.go +++ b/services/core/internal/store/creation_stream_settlement_public_test.go @@ -127,14 +127,14 @@ func TestCreationStreamPublicLifetimes(t *testing.T) { handler.ServeHTTP(w, r) })) defer server.Close() - writer := executionOwner(t, db, s).Store + sessionExecution := executionOwner(t, db, s).Sessions connect := func(environment string) { t.Helper() generation := uuid.NewString() - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, environment, generation); err != nil { + if err := sessionExecution.ReplaceEnvironmentConnection(t.Context(), tenant, environment, generation); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 1, true); err != nil { + if err := sessionExecution.ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 1, true); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/credential_matrix_http_test.go b/services/core/internal/store/credential_matrix_http_test.go index 1078e395a..786eea348 100644 --- a/services/core/internal/store/credential_matrix_http_test.go +++ b/services/core/internal/store/credential_matrix_http_test.go @@ -44,8 +44,8 @@ func TestCredentialNamespaceMatrix(t *testing.T) { } // The server composition: daemon transport beside the API handler. mux := http.NewServeMux() - mux.Handle("/api/v1/agent-daemon/enroll", runtimeenrollment.EnrollmentHandler(s)) - mux.Handle("/api/v1/agent-daemon/connection", runtimeenrollment.ConnectionHandler(s, runtimegateway.NewRegistry())) + mux.Handle("/api/v1/agent-daemon/enroll", runtimeenrollment.EnrollmentHandler(fixtureSessionService(t, db))) + mux.Handle("/api/v1/agent-daemon/connection", runtimeenrollment.ConnectionHandler(fixtureSessionStore(db), runtimegateway.NewRegistry())) mux.Handle("/", handler) server := api.CanonicalPaths(mux) call := func(method, path, token, body string) *httptest.ResponseRecorder { @@ -81,7 +81,7 @@ func TestCredentialNamespaceMatrix(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(ctx, binding.Principal.TenantID, session.ID) + environment, err := fixtureSessionStore(db).GetSessionEnvironment(ctx, binding.Principal.TenantID, session.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/device_bootstrap_binding_test.go b/services/core/internal/store/device_bootstrap_binding_test.go index 975e9fb6d..d61c74381 100644 --- a/services/core/internal/store/device_bootstrap_binding_test.go +++ b/services/core/internal/store/device_bootstrap_binding_test.go @@ -32,7 +32,7 @@ func TestDeviceCredentialCarriesPersistedAllocationNode(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -40,7 +40,7 @@ func TestDeviceCredentialCarriesPersistedAllocationNode(t *testing.T) { if err != nil { t.Fatal(err) } - authenticator := runtimegateway.NewAuthenticator(s) + authenticator := runtimegateway.NewAuthenticator(sessionAdapter(s)) auth, err := authenticator.AuthenticateBearer(t.Context(), allocation.DeviceID, bearer) if err != nil || auth.RuntimeNodeID != nodeID { t.Fatalf("authenticated node=%s want=%s error=%v", auth.RuntimeNodeID, nodeID, err) @@ -48,7 +48,7 @@ func TestDeviceCredentialCarriesPersistedAllocationNode(t *testing.T) { if _, err := authenticator.AuthenticateBearer(t.Context(), allocation.DeviceID, "wrong-token"); !errors.Is(err, runtimegateway.ErrAuthBadCredential) { t.Fatal("binding bypassed credential check", err) } - if err := s.RevokeDevice(t.Context(), tenant, allocation.DeviceID); err != nil { + if err := sessionService(t, s).RevokeDevice(t.Context(), tenant, allocation.DeviceID); err != nil { t.Fatal(err) } if _, err := authenticator.AuthenticateBearer(t.Context(), allocation.DeviceID, bearer); !errors.Is(err, runtimegateway.ErrAuthUnknownDevice) { @@ -61,7 +61,7 @@ func TestDeviceCredentialCarriesPersistedAllocationNode(t *testing.T) { func TestDeviceCredentialWithoutManagedNodeRetainsPublicRouteIdentity(t *testing.T) { s, _ := testStore(t) tenant := uuid.NewString() - ordinary, err := s.CreateDevice(t.Context(), tenant, "ordinary", runtimedevice.HashCredential("ordinary-token")) + ordinary, err := sessionService(t, s).CreateDevice(t.Context(), tenant, "ordinary", runtimedevice.HashCredential("ordinary-token")) if err != nil { t.Fatal(err) } @@ -72,12 +72,12 @@ func TestDeviceCredentialWithoutManagedNodeRetainsPublicRouteIdentity(t *testing } principal := FixtureExecutorPrincipal(t, s, uuid.NewString()) _, selfhost, key := runtimeEnrollmentFixture(t, s, principal) - enrolled, err := s.EnrollRuntime(t.Context(), selfhost.ID, executorDigest(key.Token)) + enrolled, err := sessionService(t, s).EnrollRuntime(t.Context(), selfhost.ID, executorDigest(key.Token)) if err != nil { t.Fatal(err) } for _, id := range []string{ordinary.ID, allocation.DeviceID, enrolled.DeviceID} { - credential, found, err := s.GetDeviceCredential(t.Context(), id) + credential, found, err := sessionAdapter(s).GetDeviceCredential(t.Context(), id) if err != nil || !found || credential.RuntimeNodeID != "" { t.Fatalf("non-node credential acquired allocation route: found=%v node=%s error=%v", found, credential.RuntimeNodeID, err) } @@ -85,7 +85,7 @@ func TestDeviceCredentialWithoutManagedNodeRetainsPublicRouteIdentity(t *testing if err := s.RevokeExecutorCredential(t.Context(), principal, key.KeyID); err != nil { t.Fatal(err) } - if _, found, err := s.GetDeviceCredential(t.Context(), enrolled.DeviceID); err != nil || found { + if _, found, err := sessionAdapter(s).GetDeviceCredential(t.Context(), enrolled.DeviceID); err != nil || found { t.Fatal("LEFT JOIN revived revoked executor key", err) } } diff --git a/services/core/internal/store/devices.go b/services/core/internal/store/devices.go index e7866b42a..f61266d6e 100644 --- a/services/core/internal/store/devices.go +++ b/services/core/internal/store/devices.go @@ -2,10 +2,7 @@ package store import ( "context" - "encoding/hex" "errors" - "fmt" - "strings" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -13,117 +10,9 @@ import ( "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/runtimedevice" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -// CreateDevice is operator provisioning, not a tenant-facing registration API. -func (s *Store) CreateDevice(ctx context.Context, tenantID, name, credentialHash string) (sessions.ExecutionDevice, error) { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.ExecutionDevice{}, err - } - params, err := newDeviceParams(tenant, name, credentialHash) - if err != nil { - return sessions.ExecutionDevice{}, err - } - id, err := s.queries.CreateDevice(ctx, params) - if err != nil { - return sessions.ExecutionDevice{}, fmt.Errorf("create execution device: %w", err) - } - return sessions.ExecutionDevice{ID: uuid.UUID(id.Bytes).String(), Name: params.Name}, nil -} - -func newDeviceParams(tenant pgtype.UUID, name, credentialHash string) (sqlc.CreateDeviceParams, error) { - name = strings.TrimSpace(name) - digest, err := hex.DecodeString(credentialHash) - if err != nil || len(digest) != 32 || name == "" || len(name) > 256 { - return sqlc.CreateDeviceParams{}, fmt.Errorf("%w: device name and SHA-256 credential digest required", sessions.ErrInvalidInput) - } - return sqlc.CreateDeviceParams{ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, TenantID: tenant, - Name: name, CredentialHash: pgtype.Text{String: hex.EncodeToString(digest), Valid: true}}, nil -} - -// GetDeviceCredential is used only by the shared gateway's credential verifier. -// The standalone service does not assign a product WorkspaceID. -func (s *Store) GetDeviceCredential(ctx context.Context, deviceID string) (runtimedevice.Credential, bool, error) { - id, err := parseID(deviceID) - if err != nil { - return runtimedevice.Credential{}, false, nil - } - row, err := s.queries.GetDeviceCredential(ctx, id) - if errors.Is(err, pgx.ErrNoRows) { - return runtimedevice.Credential{}, false, nil - } - if err != nil { - return runtimedevice.Credential{}, false, err - } - return runtimedevice.Credential{ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name, - Type: runtimegateway.RuntimeTypeAgentDaemon, CredentialHash: row.CredentialHash, RuntimeNodeID: row.RuntimeNodeID, RuntimeAllocationID: row.RuntimeAllocationID}, true, nil -} - -func (s *Store) RevokeDevice(ctx context.Context, tenantID, deviceID string) error { - params, err := sessionpg.DeviceLookup(tenantID, deviceID) - if err != nil { - return err - } - n, err := s.queries.RevokeDevice(ctx, sqlc.RevokeDeviceParams(params)) - if err == nil && n == 0 { - return sessions.ErrNotFound - } - return err -} - -// BindSessionDevice keeps retries stable and refuses silent filesystem moves. -// Dispatchers must obtain the full Session binding before delivery. -func (s *Store) BindSessionDevice(ctx context.Context, tenantID, sessionID, deviceID string) error { - params, err := sessionpg.DeviceLookup(tenantID, deviceID) - if err != nil { - return err - } - return s.withSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - if _, err := q.GetDevice(ctx, params); errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } else if err != nil { - return err - } - _, err := q.BindSessionDevice(ctx, sqlc.BindSessionDeviceParams{TenantID: params.TenantID, ID: session, ID_2: params.ID}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrDeviceBindingConflict - } - return err - }) -} - -func (s *Store) GetSessionDevice(ctx context.Context, tenantID, sessionID string) (sessions.ExecutionDevice, error) { - params, err := sessionpg.DeviceLookup(tenantID, sessionID) - if err != nil { - return sessions.ExecutionDevice{}, err - } - if err := s.requireInitializedEnvironment(ctx, params.TenantID, params.ID); err != nil { - return sessions.ExecutionDevice{}, err - } - return s.GetSessionRuntimeDevice(ctx, tenantID, sessionID) -} - -// GetSessionRuntimeDevice reports an authorized connection binding. It does not -// admit native execution or file access before Environment preparation completes. -func (s *Store) GetSessionRuntimeDevice(ctx context.Context, tenantID, sessionID string) (sessions.ExecutionDevice, error) { - params, err := sessionpg.DeviceLookup(tenantID, sessionID) - if err != nil { - return sessions.ExecutionDevice{}, err - } - row, err := s.queries.GetSessionDevice(ctx, sqlc.GetSessionDeviceParams(params)) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ExecutionDevice{}, sessions.ErrNotFound - } - if err != nil { - return sessions.ExecutionDevice{}, err - } - return executionDevice(row.ID, row.Name, row.EnvironmentID), nil -} - func (s *Store) GetSessionExecutionBinding(ctx context.Context, tenantID, sessionID string) (sessions.ExecutionBinding, error) { params, err := sessionpg.DeviceLookup(tenantID, sessionID) if err != nil { @@ -153,27 +42,3 @@ func executionDevice(id pgtype.UUID, name string, environmentID pgtype.UUID) ses } return value } - -func (s *Store) TouchRuntimeHeartbeat(ctx context.Context, deviceID string) (runtimedevice.HeartbeatStatus, error) { - id, err := parseID(deviceID) - if err != nil { - return runtimedevice.HeartbeatStatus{}, err - } - n, err := s.queries.TouchDevice(ctx, id) - return runtimedevice.HeartbeatStatus{Liveness: "online", Deleted: n == 0}, err -} - -func (s *Store) TouchAgentDaemonHeartbeat(ctx context.Context, heartbeat runtimedevice.Heartbeat) (runtimedevice.HeartbeatStatus, error) { - id, err := parseID(heartbeat.RuntimeID) - if err != nil { - return runtimedevice.HeartbeatStatus{}, err - } - n, err := s.queries.TouchAuthenticatedDevice(ctx, sqlc.TouchAuthenticatedDeviceParams{ - ID: id, CredentialHash: heartbeat.CredentialHash, - }) - return runtimedevice.HeartbeatStatus{Liveness: "online", Deleted: n == 0}, err -} - -// Live connectivity belongs to the gateway Registry. Only last-seen time is -// persisted, so a stale socket closing cannot overwrite a newer connection. -func (s *Store) MarkRuntimeOffline(context.Context, string) error { return nil } diff --git a/services/core/internal/store/devices_test.go b/services/core/internal/store/devices_test.go index c30960cc3..1e1d8f3e7 100644 --- a/services/core/internal/store/devices_test.go +++ b/services/core/internal/store/devices_test.go @@ -16,6 +16,8 @@ import ( "github.com/gorilla/websocket" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "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/runtime" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" @@ -25,7 +27,7 @@ import ( func registerTestDevice(t *testing.T, s *Store, tenant string) (sessions.ExecutionDevice, string) { t.Helper() secret := uuid.NewString() + uuid.NewString() - d, err := s.CreateDevice(context.Background(), tenant, "isolated executor", runtimedevice.HashCredential(secret)) + d, err := sessionService(t, s).CreateDevice(context.Background(), tenant, "isolated executor", runtimedevice.HashCredential(secret)) if err != nil { t.Fatal(err) } @@ -34,15 +36,25 @@ func registerTestDevice(t *testing.T, s *Store, tenant string) (sessions.Executi func TestDeviceBindingIsTenantScopedStableAndDurable(t *testing.T) { s, pool := testStore(t) - other, _ := testStore(t) ctx := context.Background() tenant, session := newTurnSession(t, s) otherTenant, otherSession := newTurnSession(t, s) a, _ := registerTestDevice(t, s, tenant) b, _ := registerTestDevice(t, s, tenant) foreign, _ := registerTestDevice(t, s, otherTenant) + // The binds run on an execution lease of their own, which closes before + // the pool does. + lease, err := pgunit.AcquireLease(ctx, pool) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = lease.Close(context.Background()) }) + execution, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) + if err != nil { + t.Fatal(err) + } for _, args := range [][3]string{{tenant, session.ID, foreign.ID}, {otherTenant, session.ID, foreign.ID}, {tenant, otherSession.ID, a.ID}} { - if err := s.BindSessionDevice(ctx, args[0], args[1], args[2]); !errors.Is(err, sessions.ErrNotFound) { + if err := execution.BindSessionDevice(ctx, args[0], args[1], args[2]); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("foreign binding: %v", err) } } @@ -52,11 +64,11 @@ func TestDeviceBindingIsTenantScopedStableAndDurable(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - st, id := s, a.ID + id := a.ID if i%2 == 0 { - st, id = other, b.ID + id = b.ID } - errs <- st.BindSessionDevice(ctx, tenant, session.ID, id) + errs <- execution.BindSessionDevice(ctx, tenant, session.ID, id) }() } wg.Wait() @@ -75,28 +87,31 @@ func TestDeviceBindingIsTenantScopedStableAndDurable(t *testing.T) { if success != 4 || conflicts != 4 { t.Fatalf("unstable binding: success=%d conflicts=%d", success, conflicts) } - winner, err := s.GetSessionDevice(ctx, tenant, session.ID) + winner, err := sessionAdapter(s).GetSessionDevice(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } - if _, err := s.GetSessionDevice(ctx, otherTenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSessionDevice(ctx, otherTenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("foreign lookup: %v", err) } + if err := lease.Close(ctx); err != nil { + t.Fatal(err) + } pool.Close() restarted, _ := testStore(t) - got, err := restarted.GetSessionDevice(ctx, tenant, session.ID) + got, err := sessionAdapter(restarted).GetSessionDevice(ctx, tenant, session.ID) if err != nil || got != winner { t.Fatalf("binding after restart: %+v %v", got, err) } - if err := restarted.RevokeDevice(ctx, otherTenant, winner.ID); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionService(t, restarted).RevokeDevice(ctx, otherTenant, winner.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("foreign revocation: %v", err) } for range 2 { - if err := restarted.RevokeDevice(ctx, tenant, winner.ID); err != nil { + if err := sessionService(t, restarted).RevokeDevice(ctx, tenant, winner.ID); err != nil { t.Fatal(err) } } - if _, err := restarted.GetSessionDevice(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(restarted).GetSessionDevice(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("revoked device remains dispatchable: %v", err) } } @@ -109,7 +124,7 @@ func TestStandaloneGatewayUsesExecutionCredentials(t *testing.T) { _, foreignSecret := registerTestDevice(t, s, uuid.NewString()) server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - handler, registry, err := runtime.NewGateway(s, wsURL) + handler, registry, err := runtime.NewGateway(sessionAdapter(s), sessionService(t, s), s, wsURL) if err != nil { t.Fatal(err) } @@ -169,7 +184,7 @@ func TestStandaloneGatewayUsesExecutionCredentials(t *testing.T) { if err != nil || current == previous || current.IsClosed() { t.Fatalf("replacement connection missing: %v", err) } - if err := s.RevokeDevice(ctx, tenant, a.ID); err != nil { + if err := sessionService(t, s).RevokeDevice(ctx, tenant, a.ID); err != nil { t.Fatal(err) } if err := second.WriteJSON(map[string]any{"type": proto.TypeHeartbeat, "payload": map[string]any{"version": "test"}}); err != nil { @@ -184,7 +199,7 @@ func TestStandaloneGatewayUsesExecutionCredentials(t *testing.T) { break } } - if _, err := runtimegateway.NewAuthenticator(s).AuthenticateBearer(ctx, a.ID, secret); !errors.Is(err, runtimegateway.ErrAuthUnknownDevice) { + if _, err := runtimegateway.NewAuthenticator(sessionAdapter(s)).AuthenticateBearer(ctx, a.ID, secret); !errors.Is(err, runtimegateway.ErrAuthUnknownDevice) { t.Fatalf("revoked credential survived: %v", err) } } diff --git a/services/core/internal/store/dispatch_test.go b/services/core/internal/store/dispatch_test.go index 0e36034d5..cb8fbe2a5 100644 --- a/services/core/internal/store/dispatch_test.go +++ b/services/core/internal/store/dispatch_test.go @@ -69,26 +69,26 @@ func newDispatchHarnessForSession(t *testing.T, configuration []byte, local bool } _ = json.Unmarshal(configuration, &snapshot) if snapshot.Environment.Type == "self_hosted" { - h.device, h.credential = enrollFixtureSession(t, s, h.tenant, h.session) + h.device, h.credential = enrollFixtureSession(t, s, db, h.tenant, h.session) secret = h.credential } else if local { - environment, getErr := s.GetSessionEnvironment(ctx, h.tenant, h.session.ID) + environment, getErr := fixtureSessionStore(db).GetSessionEnvironment(ctx, h.tenant, h.session.ID) if getErr != nil { t.Fatal(getErr) } - h.device, err = s.CreateEnvironmentDevice(ctx, h.tenant, environment.ID, "local runtime", runtimedevice.HashCredential(secret)) + h.device, err = store.FixtureEnvironmentDevice(ctx, db.pool, h.tenant, environment.ID, "local runtime", runtimedevice.HashCredential(secret)) } else { - h.device, err = s.CreateDevice(ctx, h.tenant, "isolated executor", runtimedevice.HashCredential(secret)) + h.device, err = fixtureSessionService(t, db).CreateDevice(ctx, h.tenant, "isolated executor", runtimedevice.HashCredential(secret)) } if err != nil { t.Fatal(err) } - if err = s.BindSessionDevice(ctx, h.tenant, h.session.ID, h.device.ID); err != nil { + if err = bindSessionDevice(t, db, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - server.Config.Handler, h.registry, err = runtime.NewGateway(s, wsURL) + server.Config.Handler, h.registry, err = runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), s, wsURL) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_connection_events_test.go b/services/core/internal/store/environment_connection_events_test.go index 29a7127c1..135b9c5a7 100644 --- a/services/core/internal/store/environment_connection_events_test.go +++ b/services/core/internal/store/environment_connection_events_test.go @@ -11,13 +11,14 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" + "github.com/jackc/pgx/v5/pgxpool" ) -func awaitEnvironmentConnectionState(t *testing.T, ctx context.Context, s *store.Store, tenant, environment, status string) { +func awaitEnvironmentConnectionState(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tenant, environment, status string) { t.Helper() deadline := time.Now().Add(5 * time.Second) for time.Now().Before(deadline) { - value, err := s.GetEnvironment(ctx, tenant, environment) + value, err := sessionReads(pool).GetEnvironment(ctx, tenant, environment) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_connection_recovery.go b/services/core/internal/store/environment_connection_recovery.go deleted file mode 100644 index 21f90dc98..000000000 --- a/services/core/internal/store/environment_connection_recovery.go +++ /dev/null @@ -1,50 +0,0 @@ -package store - -import ( - "context" - "errors" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -// ReconcileEnvironmentConnections runs before the new owner's connection producers start. -// It discards previous process generations and records loss of their connected transport. -func (s *Store) ReconcileEnvironmentConnections(ctx context.Context) error { - if err := s.checkExecutionAuthority(); err != nil { - return err - } - after := pgtype.UUID{Valid: true} - for { - var rows []sqlc.ListEnvironmentConnectionsRow - err := s.writer.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { - var err error - rows, err = s.queries.WithTx(tx).ListEnvironmentConnections(ctx, after) - return err - }) - if err != nil { - return err - } - if len(rows) == 0 { - return nil - } - for _, row := range rows { - err := s.withEnvironmentConnection(ctx, uuid.UUID(row.TenantID.Bytes).String(), uuid.UUID(row.ID.Bytes).String(), func(ctx context.Context, q *sqlc.Queries, current sqlc.GetSessionEnvironmentRow) error { - if err := q.DeleteEnvironmentConnection(ctx, current.Environment.ID); err != nil { - return err - } - if current.Environment.Status == "connected" { - return recordEnvironmentConnection(ctx, q, current, "disconnected") - } - return nil - }) - if err != nil && !errors.Is(err, sessions.ErrNotFound) { - return err - } - after = row.ID - } - } -} diff --git a/services/core/internal/store/environment_connection_recovery_test.go b/services/core/internal/store/environment_connection_recovery_test.go index f5ead182d..bfe0d6d49 100644 --- a/services/core/internal/store/environment_connection_recovery_test.go +++ b/services/core/internal/store/environment_connection_recovery_test.go @@ -21,10 +21,10 @@ func TestEnvironmentConnectionRecoveryFencesLostOwnerAcrossPages(t *testing.T) { for range 33 { tenant, session, environment := connectionFixture(t, s) generation := uuid.NewString() - if err := old.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { + if err := sessionExecution(t, old.lease).ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { t.Fatal(err) } - if err := old.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err != nil { + if err := sessionExecution(t, old.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err != nil { t.Fatal(err) } targets = append(targets, target{tenant, session, environment, generation}) @@ -33,19 +33,16 @@ func TestEnvironmentConnectionRecoveryFencesLostOwnerAcrossPages(t *testing.T) { if err := pool.QueryRow(t.Context(), "SELECT pg_terminate_backend($1,1000)", executionOwnerPID(t, pool)).Scan(&killed); err != nil || !killed { t.Fatal(killed, err) } - next := executionWriter(t, s) + next := sessionExecution(t, executionWriter(t, s).lease) first := targets[0] - if err := old.ObserveEnvironmentConnection(t.Context(), first.tenant, first.environment.ID, first.generation, 2, false); err == nil { + if err := sessionExecution(t, old.lease).ObserveEnvironmentConnection(t.Context(), first.tenant, first.environment.ID, first.generation, 2, false); err == nil { t.Fatal("lost owner wrote state") } - if err := s.ReconcileEnvironmentConnections(t.Context()); err == nil { - t.Fatal("unleased reconciliation accepted") - } if err := next.ReconcileEnvironmentConnections(t.Context()); err != nil { t.Fatal(err) } for _, target := range targets { - got, err := s.GetEnvironment(t.Context(), target.tenant, target.environment.ID) + got, err := sessionAdapter(s).GetEnvironment(t.Context(), target.tenant, target.environment.ID) if err != nil || got.Status != "disconnected" { t.Fatal("old process remained connected", got, err) } @@ -74,7 +71,7 @@ func TestEnvironmentConnectionRecoveryFencesLostOwnerAcrossPages(t *testing.T) { if err := next.ObserveEnvironmentConnection(t.Context(), first.tenant, first.environment.ID, generation, 1, true); err != nil { t.Fatal(err) } - got, err := s.GetEnvironment(t.Context(), first.tenant, first.environment.ID) + got, err := sessionAdapter(s).GetEnvironment(t.Context(), first.tenant, first.environment.ID) if err != nil || got.Status != "connected" { t.Fatal("new generation could not connect", got, err) } diff --git a/services/core/internal/store/environment_connection_worker_test.go b/services/core/internal/store/environment_connection_worker_test.go index 4313ba793..0f9a710f3 100644 --- a/services/core/internal/store/environment_connection_worker_test.go +++ b/services/core/internal/store/environment_connection_worker_test.go @@ -21,17 +21,16 @@ func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } owner := executionOwner(t, db, s) - writer := owner.Store generation := uuid.NewString() - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { + if err := owner.Sessions.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err != nil { + if err := owner.Sessions.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err != nil { t.Fatal(err) } awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, db.pool) @@ -52,7 +51,7 @@ func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { t.Error("worker cleanup did not exit") } }) - awaitEnvironmentConnectionState(t, ctx, s, tenant, environment.ID, "disconnected") + awaitEnvironmentConnectionState(t, ctx, db.pool, tenant, environment.ID, "disconnected") cancel() select { @@ -63,7 +62,7 @@ func TestEnvironmentConnectionWorkerReconcilesAndReleasesLease(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("worker did not close") } - awaitEnvironmentConnectionState(t, t.Context(), s, tenant, environment.ID, "disconnected") + awaitEnvironmentConnectionState(t, t.Context(), db.pool, tenant, environment.ID, "disconnected") if err := worker.CheckOwnership(t.Context()); err == nil { t.Fatal("worker retained lease") } diff --git a/services/core/internal/store/environment_connections.go b/services/core/internal/store/environment_connections.go index ed2b4f247..6b95a4e72 100644 --- a/services/core/internal/store/environment_connections.go +++ b/services/core/internal/store/environment_connections.go @@ -1,121 +1,10 @@ package store import ( - "context" - "errors" - - "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/google/uuid" - "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" ) -// ReplaceEnvironmentConnection starts an ordered generation, without claiming a connection. -// The producer serializes replacements and never retries an older replacement after its successor. -func (s *Store) ReplaceEnvironmentConnection(ctx context.Context, tenant, environment, generation string) error { - gen, err := parseConnectionGeneration(generation) - if err != nil { - return err - } - return s.withEnvironmentConnection(ctx, tenant, environment, func(ctx context.Context, q *sqlc.Queries, row sqlc.GetSessionEnvironmentRow) error { - if row.Environment.Status == "failed" || row.Environment.Status == "expired" { - return sessions.ErrInvalidInput - } - old, err := q.GetEnvironmentConnection(ctx, row.Environment.ID) - if err == nil && old.Generation == gen { - return nil - } - if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return err - } - if err := q.ReplaceEnvironmentConnection(ctx, sqlc.ReplaceEnvironmentConnectionParams{EnvironmentID: row.Environment.ID, Generation: gen}); err != nil { - return err - } - if row.Environment.Status == "connected" { - return recordEnvironmentConnection(ctx, q, row, "disconnected") - } - return nil - }) -} - -// ObserveEnvironmentConnection commits only a newer observation from the current generation. -// Connectivity is transport evidence, not native preparation readiness or process quiescence. -func (s *Store) ObserveEnvironmentConnection(ctx context.Context, tenant, environment, generation string, revision int64, connected bool) error { - gen, err := parseConnectionGeneration(generation) - if err != nil || revision <= 0 { - return sessions.ErrInvalidInput - } - return s.withEnvironmentConnection(ctx, tenant, environment, func(ctx context.Context, q *sqlc.Queries, row sqlc.GetSessionEnvironmentRow) error { - current, err := q.GetEnvironmentConnection(ctx, row.Environment.ID) - if errors.Is(err, pgx.ErrNoRows) { - return nil - } - if err != nil { - return err - } - if current.Generation != gen || current.Revision >= revision { - return nil - } - if row.Environment.Status == "failed" || row.Environment.Status == "expired" { - return sessions.ErrInvalidInput - } - if err := q.AdvanceEnvironmentConnection(ctx, sqlc.AdvanceEnvironmentConnectionParams{EnvironmentID: row.Environment.ID, Revision: revision}); err != nil { - return err - } - status := "disconnected" - if connected { - status = "connected" - } - if row.Environment.Status == status { - return nil - } - return recordEnvironmentConnection(ctx, q, row, status) - }) -} - -func (s *Store) withEnvironmentConnection(ctx context.Context, tenant, environment string, apply func(context.Context, *sqlc.Queries, sqlc.GetSessionEnvironmentRow) error) error { - if err := s.checkExecutionAuthority(); err != nil { - return err - } - ctx, cancel := context.WithTimeout(ctx, pgunit.ExecutionTimeout) - defer cancel() - owned, err := s.GetEnvironment(ctx, tenant, environment) - if err != nil { - return err - } - tenantID, err := parseID(tenant) - if err != nil { - return err - } - return s.withPublicSession(ctx, tenant, owned.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - row, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: tenantID, ID: session}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - return sessions.TrackInputActivity(ctx, sessionpg.BindSession(q, tenantID, session), func(ctx context.Context) error { return apply(ctx, q, row) }) - }) -} - -func recordEnvironmentConnection(ctx context.Context, q *sqlc.Queries, row sqlc.GetSessionEnvironmentRow, status string) error { - kind, err := sessions.EnvironmentType(row.Configuration) - if err != nil { - return err - } - if status != "connected" && status != "disconnected" { - return sessions.ErrInvalidInput - } - if err := q.SetEnvironmentConnectionStatus(ctx, sqlc.SetEnvironmentConnectionStatusParams{ID: row.Environment.ID, Status: status}); err != nil { - return err - } - return sessionpg.AppendChanges(ctx, q, row.Environment.SessionID, sessions.EnvironmentStateChange(uuid.UUID(row.Environment.ID.Bytes).String(), kind, status)) -} - func parseConnectionGeneration(value string) (pgtype.UUID, error) { id, err := parseID(value) if err != nil || id.Bytes == [16]byte{} { diff --git a/services/core/internal/store/environment_connections_test.go b/services/core/internal/store/environment_connections_test.go index 369b740e9..5c8fe9407 100644 --- a/services/core/internal/store/environment_connections_test.go +++ b/services/core/internal/store/environment_connections_test.go @@ -17,7 +17,7 @@ import ( func connectionFixture(t *testing.T, s *Store) (string, sessions.Session, sessions.Environment) { t.Helper() tenant, session := environmentInputSession(t, s) - environment, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -57,13 +57,13 @@ func TestEnvironmentConnectionOrdersGenerationsAndImmutableEvents(t *testing.T) first, second := uuid.NewString(), uuid.NewString() replace := func(gen string) { t.Helper() - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, gen); err != nil { + if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, gen); err != nil { t.Fatal(err) } } observe := func(gen string, rev int64, connected bool) { t.Helper() - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, gen, rev, connected); err != nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, gen, rev, connected); err != nil { t.Fatal(err) } } @@ -120,7 +120,7 @@ func TestEnvironmentConnectionOrdersGenerationsAndImmutableEvents(t *testing.T) if !reflect.DeepEqual(got, want) { t.Fatal(got, want) } - retained, err := s.GetEnvironment(t.Context(), tenant, environment.ID) + retained, err := sessionAdapter(s).GetEnvironment(t.Context(), tenant, environment.ID) if err != nil || retained.Status != "connected" { t.Fatal(retained, err) } @@ -130,20 +130,14 @@ func TestEnvironmentConnectionRequiresOwnerAndRollsBackWithEvent(t *testing.T) { s, pool := testStore(t) tenant, session, environment := connectionFixture(t, s) generation := uuid.NewString() - if err := s.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err == nil { - t.Fatal("unleased replacement accepted") - } - if err := s.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err == nil { - t.Fatal("unleased observation accepted") - } writer := executionWriter(t, s) - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { + if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { t.Fatal(err) } - if err := writer.ReplaceEnvironmentConnection(t.Context(), uuid.NewString(), environment.ID, generation); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), uuid.NewString(), environment.ID, generation); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign generation accepted", err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, session.ID, generation, 1, true); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, session.ID, generation, 1, true); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("Session ID used as Environment", err) } before := connectionSnapshot(t, pool, environment.ID) @@ -157,7 +151,7 @@ func TestEnvironmentConnectionRequiresOwnerAndRollsBackWithEvent(t *testing.T) { t.Error(err) } }) - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err == nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err == nil { t.Fatal("event failure did not abort transaction") } if after := connectionSnapshot(t, pool, environment.ID); after != before { @@ -166,14 +160,14 @@ func TestEnvironmentConnectionRequiresOwnerAndRollsBackWithEvent(t *testing.T) { if _, err := pool.Exec(t.Context(), "ALTER TABLE session_events DROP CONSTRAINT "+constraint); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err != nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true); err != nil { t.Fatal(err) } if err := writer.lease.Close(t.Context()); err != nil { t.Fatal(err) } before = connectionSnapshot(t, pool, environment.ID) - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 2, false); err == nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 2, false); err == nil { t.Fatal("closed owner used pooled writer") } if after := connectionSnapshot(t, pool, environment.ID); after != before { @@ -188,7 +182,7 @@ func TestEnvironmentConnectionDoesNotReviveDeletedOrTerminalResources(t *testing tenant, session, environment := connectionFixture(t, s) writer := executionWriter(t, s) generation := uuid.NewString() - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { + if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, generation); err != nil { t.Fatal(err) } if status == "deleted" { @@ -203,10 +197,10 @@ func TestEnvironmentConnectionDoesNotReviveDeletedOrTerminalResources(t *testing before := connectionSnapshot(t, pool, environment.ID) for _, operation := range []func() error{ func() error { - return writer.ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, uuid.NewString()) + return sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, environment.ID, uuid.NewString()) }, func() error { - return writer.ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true) + return sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment.ID, generation, 1, true) }, } { if err := operation(); err == nil { diff --git a/services/core/internal/store/environment_device_test.go b/services/core/internal/store/environment_device_test.go index 9a7a116ee..c9578f93d 100644 --- a/services/core/internal/store/environment_device_test.go +++ b/services/core/internal/store/environment_device_test.go @@ -34,7 +34,7 @@ func TestWorkerEnvironmentSelectsCapableDeviceWithoutMovingBinding(t *testing.T) t.Fatal("incapable enrolled device received work", frame.Type) case <-time.After(time.Second): } - if _, err := h.s.GetSessionDevice(t.Context(), h.tenant, pending.SessionID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := fixtureSessionStore(h.db).GetSessionDevice(t.Context(), h.tenant, pending.SessionID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unregistered Runtime was assigned general compute", err) } session, err := h.s.GetSession(t.Context(), h.tenant, pending.SessionID) @@ -44,11 +44,11 @@ func TestWorkerEnvironmentSelectsCapableDeviceWithoutMovingBinding(t *testing.T) other := connectFixtureRuntime(t, h, session) otherFrames := workerFrames(t, other) request := nextWorkerFrame(t, otherFrames, proto.TypeExecutionPrepare) - selected, err := h.s.GetSessionDevice(t.Context(), h.tenant, pending.SessionID) + selected, err := fixtureSessionStore(h.db).GetSessionDevice(t.Context(), h.tenant, pending.SessionID) if err != nil || selected.ID != other.device.ID { t.Fatal("enrollment did not retain exact Runtime", err) } - original, err := h.s.GetSessionDevice(t.Context(), h.tenant, bound.SessionID) + original, err := fixtureSessionStore(h.db).GetSessionDevice(t.Context(), h.tenant, bound.SessionID) if err != nil || original.ID != originalRuntime.device.ID { t.Fatal("existing binding moved to a capable Runtime", err) } diff --git a/services/core/internal/store/environment_directory_test.go b/services/core/internal/store/environment_directory_test.go index 537b5be6e..087c97f60 100644 --- a/services/core/internal/store/environment_directory_test.go +++ b/services/core/internal/store/environment_directory_test.go @@ -21,7 +21,7 @@ type directoryResult struct { func directoryWorker(t *testing.T) (*dispatchHarness, *execution.Worker, sessions.Environment) { t.Helper() h := newDispatchHarnessForSession(t, []byte(`{"agent":{"model":"unavailable-model"},"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`), true) - environment, err := h.s.GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) + environment, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_executor_command_test.go b/services/core/internal/store/environment_executor_command_test.go index 98b0228cf..3dc6175ca 100644 --- a/services/core/internal/store/environment_executor_command_test.go +++ b/services/core/internal/store/environment_executor_command_test.go @@ -53,7 +53,7 @@ func TestEnvironmentExecutorOperatorCommand(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionReads(pool).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -65,14 +65,14 @@ func TestEnvironmentExecutorOperatorCommand(t *testing.T) { if next == first { t.Fatal("rotation returned the same key") } - if _, err := s.AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(first)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionReads(pool).AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(first)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("old command credential retained authority", err) } - if owner, err := s.AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(next)); err != nil || owner != tenant { + if owner, err := sessionReads(pool).AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(next)); err != nil || owner != tenant { t.Fatal("rotated command credential failed", err) } command(tenant, true, "--revoke") - if _, err := s.AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(next)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionReads(pool).AuthenticateEnvironmentExecutor(t.Context(), environment.ID, runtimedevice.HashCredential(next)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("revoked command credential retained authority", err) } t.Log("built operator command issued before Session creation, rejected duplicate/foreign requests, rotated and revoked durable credentials") diff --git a/services/core/internal/store/environment_executor_credentials.go b/services/core/internal/store/environment_executor_credentials.go index 890d2a045..d9726abd2 100644 --- a/services/core/internal/store/environment_executor_credentials.go +++ b/services/core/internal/store/environment_executor_credentials.go @@ -2,15 +2,11 @@ package store import ( "context" - "crypto/sha256" - "encoding/hex" "errors" - "fmt" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" ) @@ -138,23 +134,3 @@ func (s *Store) executorCredentialRestriction(ctx context.Context, principal ide } return restriction, err } - -// AuthenticateEnvironmentExecutor checks the current key and recorded Session creator in one database snapshot. -func (s *Store) AuthenticateEnvironmentExecutor(ctx context.Context, environment, digest string) (string, error) { - id, err := parseID(environment) - if err != nil { - return "", sessions.ErrNotFound - } - hash, err := hex.DecodeString(digest) - if err != nil || len(hash) != sha256.Size { - return "", sessions.ErrNotFound - } - tenant, err := s.queries.AuthenticateEnvironmentExecutor(ctx, sqlc.AuthenticateEnvironmentExecutorParams{EnvironmentID: id, TokenSha256: hex.EncodeToString(hash)}) - if errors.Is(err, pgx.ErrNoRows) { - return "", sessions.ErrNotFound - } - if err != nil { - return "", fmt.Errorf("authenticate environment executor: %w", err) - } - return uuid.UUID(tenant.Bytes).String(), nil -} diff --git a/services/core/internal/store/environment_executor_credentials_test.go b/services/core/internal/store/environment_executor_credentials_test.go index 6f9563354..af9ad84fa 100644 --- a/services/core/internal/store/environment_executor_credentials_test.go +++ b/services/core/internal/store/environment_executor_credentials_test.go @@ -26,7 +26,7 @@ func TestEnvironmentExecutorCredentialLifecycle(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(ctx, tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } @@ -43,7 +43,7 @@ func TestEnvironmentExecutorCredentialLifecycle(t *testing.T) { } check := func(st *Store, token string, allowed bool) { t.Helper() - owner, err := st.AuthenticateEnvironmentExecutor(ctx, environment.ID, executorDigest(token)) + owner, err := sessionAdapter(st).AuthenticateEnvironmentExecutor(ctx, environment.ID, executorDigest(token)) if allowed { if err != nil || owner != tenant { t.Fatal("credential not accepted for owner", err) @@ -54,7 +54,7 @@ func TestEnvironmentExecutorCredentialLifecycle(t *testing.T) { } check(s, token, true) check(s, "caller/device/harness/grant", false) - if _, err := s.AuthenticateEnvironmentExecutor(ctx, uuid.NewString(), executorDigest(token)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, uuid.NewString(), executorDigest(token)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign Environment accepted", err) } if _, err := s.IssueExecutorCredential(ctx, principal, environment.ID, environment.ID); !errors.Is(err, sessions.ErrExecutorCredentialExists) { @@ -119,7 +119,7 @@ func TestEnvironmentExecutorConcurrentIssueAndDeletion(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(ctx, tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } @@ -167,11 +167,11 @@ func TestEnvironmentExecutorConcurrentIssueAndDeletion(t *testing.T) { wg.Wait() close(rotated) for token := range rotated { - if _, err := s.AuthenticateEnvironmentExecutor(ctx, environment.ID, executorDigest(token)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, environment.ID, executorDigest(token)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("racing rotation authorized deleted Environment", err) } } - if _, err := s.AuthenticateEnvironmentExecutor(ctx, environment.ID, executorDigest(original)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, environment.ID, executorDigest(original)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("original key survived deletion", err) } } diff --git a/services/core/internal/store/environment_executor_management.go b/services/core/internal/store/environment_executor_management.go index 6a3be82a9..361915520 100644 --- a/services/core/internal/store/environment_executor_management.go +++ b/services/core/internal/store/environment_executor_management.go @@ -9,6 +9,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" "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/projects" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" @@ -177,7 +178,7 @@ func (s *Store) selfHostedExecutorTarget(ctx context.Context, principal identity if err := principal.Validate(); err != nil { return sessions.ErrInvalidInput } - owned, err := s.GetEnvironment(ctx, principal.TenantID, environment) + owned, err := sessionpg.LoadEnvironment(ctx, s.queries, principal.TenantID, environment) if err != nil { return err } diff --git a/services/core/internal/store/environment_executor_management_test.go b/services/core/internal/store/environment_executor_management_test.go index df31ee905..2e87de334 100644 --- a/services/core/internal/store/environment_executor_management_test.go +++ b/services/core/internal/store/environment_executor_management_test.go @@ -30,7 +30,7 @@ func TestProjectEnvironmentExecutorManagement(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(ctx, p.TenantID, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(ctx, p.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -49,7 +49,7 @@ func TestProjectEnvironmentExecutorManagement(t *testing.T) { if err != nil || issued.KeyID != keyID || issued.EnvironmentID != one.ID || issued.Token == "" { t.Fatal("issue", err) } - if _, err := s.AuthenticateEnvironmentExecutor(ctx, one.ID, executorDigest(issued.Token)); err != nil { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, one.ID, executorDigest(issued.Token)); err != nil { t.Fatal("issued credential does not authenticate", err) } // A lost issuance response must not cause a new key or replace the old secret. @@ -90,7 +90,7 @@ func TestProjectEnvironmentExecutorManagement(t *testing.T) { if err != nil || rotated.Token == issued.Token { t.Fatal("rotation", err) } - if _, err := s.AuthenticateEnvironmentExecutor(ctx, one.ID, executorDigest(issued.Token)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, one.ID, executorDigest(issued.Token)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("old secret", err) } for range 2 { @@ -98,7 +98,7 @@ func TestProjectEnvironmentExecutorManagement(t *testing.T) { t.Fatal(err) } } - if _, err := s.AuthenticateEnvironmentExecutor(ctx, one.ID, executorDigest(rotated.Token)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, one.ID, executorDigest(rotated.Token)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("revoked secret", err) } if listed, err := s.ListProjectExecutorCredentials(ctx, p, one.ID); err != nil || len(listed) != 1 || listed[0].RevokedAt == nil { @@ -153,7 +153,7 @@ func TestArchivedProjectExecutorCredentials(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(ctx, p.TenantID, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(ctx, p.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -202,7 +202,7 @@ func TestProjectExecutorConnectionState(t *testing.T) { if err != nil || state.Connection.DeviceID != "" || state.Connection.EnrolledAt != nil || state.Connection.BoundKeyID != nil || state.Connection.LastSeenAt != nil { t.Fatal("never enrolled", state, err) } - enrolled, err := s.EnrollRuntime(ctx, env.ID, executorDigest(key.Token)) + enrolled, err := sessionService(t, s).EnrollRuntime(ctx, env.ID, executorDigest(key.Token)) if err != nil { t.Fatal(err) } @@ -250,7 +250,7 @@ func TestProjectExecutorConnectionState(t *testing.T) { if _, err = pool.Exec(ctx, "UPDATE environments SET status='expired' WHERE id=$1", env.ID); err != nil { t.Fatal(err) } - if _, err = s.AuthenticateEnvironmentExecutor(ctx, env.ID, executorDigest(rotated.Token)); err != nil { + if _, err = sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, env.ID, executorDigest(rotated.Token)); err != nil { t.Fatal("fixture executor key should still authenticate", err) } state = check() diff --git a/services/core/internal/store/environment_file_writes.go b/services/core/internal/store/environment_file_writes.go index 0c56b14f3..f603cdfa9 100644 --- a/services/core/internal/store/environment_file_writes.go +++ b/services/core/internal/store/environment_file_writes.go @@ -25,7 +25,7 @@ func (s *Store) ReserveEnvironmentFileWrite(ctx context.Context, tenant, environ if !key.Valid() { return sessions.EnvironmentFileWrite{}, sessions.ErrInvalidInput } - owned, err := s.GetEnvironment(ctx, tenant, environment) + owned, err := sessionpg.LoadEnvironment(ctx, s.queries, tenant, environment) if err != nil { return sessions.EnvironmentFileWrite{}, err } diff --git a/services/core/internal/store/environment_file_writes_test.go b/services/core/internal/store/environment_file_writes_test.go index 96021d168..51b5e7395 100644 --- a/services/core/internal/store/environment_file_writes_test.go +++ b/services/core/internal/store/environment_file_writes_test.go @@ -21,11 +21,11 @@ type fileWriteFixture struct { func newFileWriteFixture(t *testing.T) fileWriteFixture { t.Helper() - s, _ := testStore(t) + s, pool := testStore(t) writer := executionWriter(t, s) tenant := uuid.NewString() session, env := localEnvironment(t, s, tenant) - host, err := s.CreateEnvironmentDevice(t.Context(), tenant, env.ID, "file owner", runtimedevice.HashCredential(uuid.NewString())) + host, err := FixtureEnvironmentDevice(t.Context(), pool, tenant, env.ID, "file owner", runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } @@ -64,7 +64,7 @@ func TestEnvironmentFileWriteRetainsUnknownAcrossLeaseLoss(t *testing.T) { if _, err := reopened.SubmitMessage(ctx, f.tenant, f.session.ID, "direct", messageInput("new").Payload); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("direct admission bypassed write", err) } - if _, err := reopened.GetEnvironment(ctx, f.tenant, f.env.ID); err != nil { + if _, err := sessionAdapter(reopened).GetEnvironment(ctx, f.tenant, f.env.ID); err != nil { t.Fatal("write gate prevented metadata read", err) } if _, err := reopened.GetSession(ctx, f.tenant, f.session.ID); err != nil { diff --git a/services/core/internal/store/environment_initial_input_test.go b/services/core/internal/store/environment_initial_input_test.go index 1762ff5f3..1a9b2badf 100644 --- a/services/core/internal/store/environment_initial_input_test.go +++ b/services/core/internal/store/environment_initial_input_test.go @@ -144,10 +144,10 @@ func TestEnvironmentInitialInputCreationRetainsCursorIdentityAndPromotion(t *tes } writer := executionWriter(t, s) generation := uuid.NewString() - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, session.Environment.ID, generation); err != nil { + if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, session.Environment.ID, generation); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, session.Environment.ID, generation, 1, true); err != nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, session.Environment.ID, generation, 1, true); err != nil { t.Fatal(err) } connectedStatus := "idle" @@ -244,10 +244,10 @@ func TestEnvironmentInitialInputExpiryHasNoTurnAndCannotReplay(t *testing.T) { t.Fatal("reopened creation retry reset expiry", got) } generation := uuid.NewString() - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, session.Environment.ID, generation); err != nil { + if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, session.Environment.ID, generation); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, session.Environment.ID, generation, 1, true); err != nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, session.Environment.ID, generation, 1, true); err != nil { t.Fatal(err) } late, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) diff --git a/services/core/internal/store/environment_initialization.go b/services/core/internal/store/environment_initialization.go deleted file mode 100644 index 6bc15cbaf..000000000 --- a/services/core/internal/store/environment_initialization.go +++ /dev/null @@ -1,116 +0,0 @@ -package store - -import ( - "context" - "errors" - - "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" -) - -func (s *Store) ListEnvironmentInitializations(ctx context.Context, after string) ([]sessions.EnvironmentInitialization, 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.ListEnvironmentInitializations(ctx, id) - if err != nil { - return nil, err - } - result := make([]sessions.EnvironmentInitialization, 0, len(rows)) - for _, row := range rows { - result = append(result, sessions.EnvironmentInitialization{EnvironmentID: runtimeUUID(row.ID), SessionID: runtimeUUID(row.SessionID), TenantID: runtimeUUID(row.TenantID), DeviceID: runtimeUUID(row.DeviceID), State: row.Initialization, Engine: row.Engine}) - } - return result, nil -} - -func (s *Store) mutateEnvironmentInitialization(ctx context.Context, owner sessions.EnvironmentInitialization, apply func(*sqlc.Queries, sqlc.GetSessionEnvironmentRow) error) error { - if err := s.checkExecutionAuthority(); err != nil { - return err - } - tenant, err := parseID(owner.TenantID) - if err != nil { - return err - } - return s.withPublicSession(ctx, owner.TenantID, owner.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - row, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: tenant, ID: session}) - if err != nil { - return err - } - if runtimeUUID(row.Environment.ID) != owner.EnvironmentID { - return sessions.ErrDeviceBindingConflict - } - if row.Environment.Status == "failed" || row.Environment.Status == "expired" { - return sessions.ErrNotFound - } - return apply(q, row) - }) -} - -func (s *Store) ClaimEnvironmentInitialization(ctx context.Context, owner sessions.EnvironmentInitialization) error { - return s.mutateEnvironmentInitialization(ctx, owner, func(q *sqlc.Queries, row sqlc.GetSessionEnvironmentRow) error { - if err := checkInitializationDevice(ctx, q, owner, row); err != nil { - return err - } - count, err := q.ClaimEnvironmentInitialization(ctx, row.Environment.ID) - if err == nil && count != 1 { - return sessions.ErrTurnConflict - } - return err - }) -} - -func checkInitializationDevice(ctx context.Context, q *sqlc.Queries, owner sessions.EnvironmentInitialization, row sqlc.GetSessionEnvironmentRow) error { - tenant, _ := parseID(owner.TenantID) - bound, err := q.GetSessionDevice(ctx, sqlc.GetSessionDeviceParams{TenantID: tenant, ID: row.Environment.SessionID}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - if runtimeUUID(bound.ID) != owner.DeviceID || bound.EnvironmentID != row.Environment.ID { - return sessions.ErrTurnConflict - } - return nil -} - -func (s *Store) CompleteEnvironmentInitialization(ctx context.Context, owner sessions.EnvironmentInitialization) error { - return s.mutateEnvironmentInitialization(ctx, owner, func(q *sqlc.Queries, row sqlc.GetSessionEnvironmentRow) error { - if err := checkInitializationDevice(ctx, q, owner, row); err != nil { - return err - } - count, err := q.CompleteEnvironmentInitialization(ctx, row.Environment.ID) - if err == nil && count != 1 { - return sessions.ErrTurnConflict - } - return err - }) -} - -// Failure records state and settles work, never destroys compute or a workspace. -func (s *Store) FailEnvironmentInitialization(ctx context.Context, owner sessions.EnvironmentInitialization, failure sessions.ProvisioningFailure) error { - return s.mutateEnvironmentInitialization(ctx, owner, func(q *sqlc.Queries, row sqlc.GetSessionEnvironmentRow) error { - if row.Environment.Initialization == "complete" { - return sessions.ErrTurnConflict - } - if err := q.FailEnvironmentInitialization(ctx, row.Environment.ID); err != nil { - return err - } - environment, err := sessionpg.EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, nil) - if err != nil { - return err - } - return sessions.FailEnvironment(ctx, sessionpg.BindSession(q, row.TenantID, row.Environment.SessionID), environment, failure.Reason(), failure.Detail()) - }) -} diff --git a/services/core/internal/store/environment_initialization_test.go b/services/core/internal/store/environment_initialization_test.go index ffd7e8d9f..66b8554f4 100644 --- a/services/core/internal/store/environment_initialization_test.go +++ b/services/core/internal/store/environment_initialization_test.go @@ -5,6 +5,13 @@ import ( "context" "encoding/json" "errors" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + "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/execution" @@ -12,32 +19,27 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/google/uuid" - "net/http" - "net/http/httptest" - "strings" - "sync" - "testing" - "time" + "github.com/jackc/pgx/v5/pgxpool" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) -func initializationState(t *testing.T, s *store.Store, tenant, environment string) string { +func initializationState(t *testing.T, pool *pgxpool.Pool, tenant, environment string) string { t.Helper() - value, err := s.GetEnvironment(t.Context(), tenant, environment) + value, err := sessionReads(pool).GetEnvironment(t.Context(), tenant, environment) if err != nil { t.Fatal(err) } return value.Initialization } -func awaitInitialization(t *testing.T, s *store.Store, tenant, environment, state string) { +func awaitInitialization(t *testing.T, pool *pgxpool.Pool, tenant, environment, state string) { t.Helper() deadline := time.Now().Add(10 * time.Second) for time.Now().Before(deadline) { - if initializationState(t, s, tenant, environment) == state { + if initializationState(t, pool, tenant, environment) == state { return } time.Sleep(10 * time.Millisecond) @@ -64,7 +66,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) + environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -72,12 +74,12 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test if err != nil { t.Fatal(err) } - enrolled, err := s.EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) + enrolled, err := fixtureSessionService(t, db).EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) if err != nil { t.Fatal(err) } registry := runtimegateway.NewRegistry() - handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(s), Registry: registry}) + handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(fixtureSessionStore(db)), Registry: registry}) server := httptest.NewServer(http.HandlerFunc(handler.WS)) defer server.Close() worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: registry}) @@ -109,7 +111,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test actions = append(actions, action) mu.Unlock() if outcome == "revoked" { - if err := s.RevokeDevice(t.Context(), principal.TenantID, enrolled.DeviceID); err != nil { + if err := fixtureSessionService(t, db).RevokeDevice(t.Context(), principal.TenantID, enrolled.DeviceID); err != nil { t.Error(err) } return completedInitialization(request, data) @@ -120,7 +122,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test return completedInitialization(request, data) } time.Sleep(350 * time.Millisecond) - if initializationState(t, s, principal.TenantID, environment.ID) != "pending" { + if initializationState(t, db.pool, principal.TenantID, environment.ID) != "pending" { t.Fatal("unconnected preparation was consumed") } bootstrap := sandbox.Bootstrap{DeviceID: enrolled.DeviceID, Credential: key.Token} @@ -131,7 +133,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test if outcome != "completed" { want = "failed" } - awaitInitialization(t, s, principal.TenantID, environment.ID, want) + awaitInitialization(t, db.pool, principal.TenantID, environment.ID, want) if _, err := s.GetRuntimeAllocation(t.Context(), principal.TenantID, environment.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("self-hosted preparation fabricated allocation", err) } @@ -181,7 +183,7 @@ func TestEnvironmentInitializationRevocationBeforeClaim(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) + environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -189,15 +191,15 @@ func TestEnvironmentInitializationRevocationBeforeClaim(t *testing.T) { if err != nil { t.Fatal(err) } - enrolled, err := s.EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) + enrolled, err := fixtureSessionService(t, db).EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) if err != nil { t.Fatal(err) } return sessions.EnvironmentInitialization{EnvironmentID: environment.ID, SessionID: session.ID, TenantID: principal.TenantID, DeviceID: enrolled.DeviceID, State: "pending", Engine: "codex"} } revoked, other := create(), create() - owned := executionOwner(t, db, s).Store - if err := s.RevokeDevice(t.Context(), principal.TenantID, revoked.DeviceID); err != nil { + owned := executionOwner(t, db, s).Sessions + if err := fixtureSessionService(t, db).RevokeDevice(t.Context(), principal.TenantID, revoked.DeviceID); err != nil { t.Fatal(err) } if err := owned.ClaimEnvironmentInitialization(t.Context(), revoked); !errors.Is(err, sessions.ErrNotFound) { @@ -208,7 +210,7 @@ func TestEnvironmentInitializationRevocationBeforeClaim(t *testing.T) { if err := owned.ClaimEnvironmentInitialization(t.Context(), stale); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatalf("stale binding escaped normal admission handling: %v", err) } - if initializationState(t, s, other.TenantID, other.EnvironmentID) != "pending" { + if initializationState(t, db.pool, other.TenantID, other.EnvironmentID) != "pending" { t.Fatal("stale claim changed preparation state") } if err := owned.ClaimEnvironmentInitialization(t.Context(), other); err != nil { diff --git a/services/core/internal/store/environment_input_activity_test.go b/services/core/internal/store/environment_input_activity_test.go index 21937bdd6..1ed22d4ec 100644 --- a/services/core/internal/store/environment_input_activity_test.go +++ b/services/core/internal/store/environment_input_activity_test.go @@ -54,10 +54,10 @@ func TestEnvironmentInputActivityWaitsBeforeTurnAndClearsOnConnection(t *testing reserveEnvironmentInput(t, s, tenant, session.ID, "waiting") writer := executionWriter(t, s) generation := uuid.NewString() - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, environment, generation); err != nil { + if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, environment, generation); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 1, true); err != nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 1, true); err != nil { t.Fatal(err) } idle := requireEnvironmentInputActivity(t, s, tenant, session.ID, "idle", "") @@ -75,14 +75,14 @@ func TestEnvironmentInputActivityWaitsBeforeTurnAndClearsOnConnection(t *testing if !idle.PendingInput || changes[2].Settled || changes[0].Settled { t.Fatal("pending input reported as settled") } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 1, false); err != nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 1, false); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 2, false); err != nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 2, false); err != nil { t.Fatal(err) } requireEnvironmentInputActivity(t, s, tenant, session.ID, "requires_action", environment) - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 3, true); err != nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 3, true); err != nil { t.Fatal(err) } if _, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID); err != nil { @@ -187,14 +187,14 @@ func TestEnvironmentInputActivityRollsBackReservationAndConnection(t *testing.T) value := requireEnvironmentInputActivity(t, s, tenant, session.ID, "requires_action", session.Environment.ID) writer := executionWriter(t, s) generation := uuid.NewString() - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, value.Environment.ID, generation); err != nil { + if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, value.Environment.ID, generation); err != nil { t.Fatal(err) } before := connectionSnapshot(t, pool, value.Environment.ID) if _, err := pool.Exec(t.Context(), "ALTER TABLE session_events ADD CONSTRAINT "+constraint+" CHECK (session_id <> '"+session.ID+"' OR payload->'event'->>'type' <> 'agent.session.idle') NOT VALID"); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(t.Context(), tenant, value.Environment.ID, generation, 1, true); err == nil { + if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, value.Environment.ID, generation, 1, true); err == nil { t.Fatal("connection committed without activity") } if connection := connectionSnapshot(t, pool, value.Environment.ID); connection != before { @@ -210,10 +210,10 @@ func TestEnvironmentInputActivityRecoversWaitingActionAndHidesDeletion(t *testin old := executionWriter(t, s) generation := uuid.NewString() environment := session.Environment.ID - if err := old.ReplaceEnvironmentConnection(t.Context(), tenant, environment, generation); err != nil { + if err := sessionExecution(t, old.lease).ReplaceEnvironmentConnection(t.Context(), tenant, environment, generation); err != nil { t.Fatal(err) } - if err := old.ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 1, true); err != nil { + if err := sessionExecution(t, old.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 1, true); err != nil { t.Fatal(err) } requireEnvironmentInputActivity(t, s, tenant, session.ID, "idle", "") @@ -223,7 +223,7 @@ func TestEnvironmentInputActivityRecoversWaitingActionAndHidesDeletion(t *testin } awaitRelease() next := executionWriter(t, s) - if err := next.ReconcileEnvironmentConnections(t.Context()); err != nil { + if err := sessionExecution(t, next.lease).ReconcileEnvironmentConnections(t.Context()); err != nil { t.Fatal(err) } requireEnvironmentInputActivity(t, s, tenant, session.ID, "requires_action", environment) @@ -235,7 +235,7 @@ func TestEnvironmentInputActivityRecoversWaitingActionAndHidesDeletion(t *testin if err != nil { t.Fatal(err) } - if err := next.ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 2, true); err != nil { + if err := sessionExecution(t, next.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 2, true); err != nil { t.Fatal(err) } if after, err := s.SessionEventCursor(t.Context(), tenant, session.ID); err != nil || after != cursor { diff --git a/services/core/internal/store/environment_input_claim_test.go b/services/core/internal/store/environment_input_claim_test.go index 5c27d3cd6..af872bf5b 100644 --- a/services/core/internal/store/environment_input_claim_test.go +++ b/services/core/internal/store/environment_input_claim_test.go @@ -1,9 +1,10 @@ package store import ( - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "sync" "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) func TestEnvironmentInputConcurrentPromotionClaimsOnce(t *testing.T) { diff --git a/services/core/internal/store/environment_installation_test.go b/services/core/internal/store/environment_installation_test.go index 4e3f3fc67..a264db7ac 100644 --- a/services/core/internal/store/environment_installation_test.go +++ b/services/core/internal/store/environment_installation_test.go @@ -30,7 +30,7 @@ func TestEnvironmentInstallationClaimLifetimeAndRetries(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(ctx, p.TenantID, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(ctx, p.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -81,7 +81,7 @@ func TestEnvironmentInstallationClaimLifetimeAndRetries(t *testing.T) { if err := s.ClaimEnvironmentInstallation(ctx, token, "build", secrets[winner]); err != nil { t.Fatal("lost-response retry", err) } - if _, err := s.AuthenticateEnvironmentExecutor(ctx, environment.ID, executorDigest(secrets[winner])); err != nil { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, environment.ID, executorDigest(secrets[winner])); err != nil { t.Fatal(err) } if err := s.RevokeExecutorCredential(ctx, p, environment.ID); err != nil { diff --git a/services/core/internal/store/environment_packages_test.go b/services/core/internal/store/environment_packages_test.go deleted file mode 100644 index 2510ad3f5..000000000 --- a/services/core/internal/store/environment_packages_test.go +++ /dev/null @@ -1,32 +0,0 @@ -package store - -import ( - "bytes" - "testing" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" - "github.com/google/uuid" -) - -func TestStoredSystemPackagesRejected(t *testing.T) { - cipher, err := credentialcrypto.New(bytes.Repeat([]byte{5}, 32)) - if err != nil { - t.Fatal(err) - } - s := NewWithCredentialCipher(nil, cipher) - tenant, session := uuid.NewString(), uuid.NewString() - for _, value := range []string{`null`, `[]`, `["jq"]`} { - packages := []byte(`{"npm":[],"python":[],"system":` + value + "}") - plain := append([]byte(`{"packages":`), packages...) - plain = append(plain, '}') - encrypted, err := cipher.SealEnvironmentSetup(plain, credentialcrypto.EnvironmentSetupBinding{TenantID: tenant, Resource: "session", OwnerID: session, Field: "initialization"}) - if err != nil { - t.Fatal(err) - } - var output environmentconfig.Setup - if err := s.openEnvironmentSetup(tenant, "session", session, "initialization", encrypted, &output); err == nil { - t.Fatal("snapshot silently ignored removed system packages", value) - } - } -} diff --git a/services/core/internal/store/environment_plugins_test.go b/services/core/internal/store/environment_plugins_test.go index 61eaebda6..241af01d6 100644 --- a/services/core/internal/store/environment_plugins_test.go +++ b/services/core/internal/store/environment_plugins_test.go @@ -56,7 +56,7 @@ func TestPluginsFrozenInSession(t *testing.T) { if err = json.Unmarshal(session.Configuration, &cfg); err != nil || len(cfg.Environment.Plugins) != 1 || !reflect.DeepEqual(cfg.Environment.Directories, setup.CapabilityDirectories) { t.Fatal("frozen public metadata", err) } - frozen, err := s.ReadEnvironmentSetup(t.Context(), tenant, session.ID) + frozen, err := sessionAdapter(s).ReadEnvironmentSetup(t.Context(), tenant, session.ID) if err != nil || !reflect.DeepEqual(frozen.Plugins, setup.Plugins) || !reflect.DeepEqual(frozen.CapabilityDirectories, setup.CapabilityDirectories) { t.Fatal("frozen snapshot changed", err) } @@ -64,7 +64,7 @@ func TestPluginsFrozenInSession(t *testing.T) { if err != nil || retry.ID != session.ID { t.Fatal("retry", err) } - if _, err = s.ReadEnvironmentSetup(t.Context(), foreign, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).ReadEnvironmentSetup(t.Context(), foreign, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign snapshot", err) } } diff --git a/services/core/internal/store/environment_retrieve_public_test.go b/services/core/internal/store/environment_retrieve_public_test.go index cc9036bec..658d013e6 100644 --- a/services/core/internal/store/environment_retrieve_public_test.go +++ b/services/core/internal/store/environment_retrieve_public_test.go @@ -76,7 +76,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { return result } result := run() - before, err := s.GetEnvironment(t.Context(), tenant, result["environment_id"]) + before, err := fixtureSessionStore(db).GetEnvironment(t.Context(), tenant, result["environment_id"]) if err != nil { t.Fatal(err) } @@ -98,7 +98,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { settings[key] = value } run() - after, err := reopened.GetEnvironment(t.Context(), tenant, before.ID) + after, err := fixtureSessionStore(reopenedDB).GetEnvironment(t.Context(), tenant, before.ID) if err != nil || !reflect.DeepEqual(before, after) { t.Fatal("public retrieval changed durable Environment state", err) } diff --git a/services/core/internal/store/environment_runtime_fixture_test.go b/services/core/internal/store/environment_runtime_fixture_test.go index c400fd523..2671ce100 100644 --- a/services/core/internal/store/environment_runtime_fixture_test.go +++ b/services/core/internal/store/environment_runtime_fixture_test.go @@ -14,9 +14,9 @@ import ( "github.com/gorilla/websocket" ) -func enrollFixtureSession(t *testing.T, s *store.Store, tenant string, session sessions.Session) (sessions.ExecutionDevice, string) { +func enrollFixtureSession(t *testing.T, s *store.Store, db fixtureDB, tenant string, session sessions.Session) (sessions.ExecutionDevice, string) { t.Helper() - environment, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -25,11 +25,11 @@ func enrollFixtureSession(t *testing.T, s *store.Store, tenant string, session s if err != nil { t.Fatal(err) } - enrolled, err := s.EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) + enrolled, err := fixtureSessionService(t, db).EnrollRuntime(t.Context(), environment.ID, runtimedevice.HashCredential(key.Token)) if err != nil || enrolled.EnvironmentID != environment.ID || enrolled.SessionID != session.ID || enrolled.WorkspaceDirectory != "/workspace" { t.Fatalf("Runtime enrollment: %+v %v", enrolled, err) } - bound, err := s.GetSessionDevice(t.Context(), tenant, session.ID) + bound, err := fixtureSessionStore(db).GetSessionDevice(t.Context(), tenant, session.ID) if err != nil || bound.ID != enrolled.DeviceID || bound.EnvironmentID != environment.ID { t.Fatalf("Runtime binding: %+v %v", bound, err) } @@ -41,7 +41,7 @@ func connectFixtureRuntime(t *testing.T, h *dispatchHarness, session sessions.Se // The Runtime shares the harness's Core, not its connection or write lock. other := &dispatchHarness{t: h.t, s: h.s, db: h.db, lease: h.lease, owned: h.owned, d: h.d, tenant: h.tenant, session: session, registry: h.registry, url: h.url, admissions: h.admissions, environments: h.environments} - other.device, other.credential = enrollFixtureSession(t, h.s, h.tenant, session) + other.device, other.credential = enrollFixtureSession(t, h.s, h.db, h.tenant, session) u, err := url.Parse(h.url) if err != nil { t.Fatal(err) @@ -75,7 +75,7 @@ func completeEmptyArtifactExport(t *testing.T, h *dispatchHarness, frames ...<-c } frame := read(proto.TypeExecutionPrepare) var prepare proto.ExecutionPreparePayload - environment, err := h.s.GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) + environment, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) if err != nil || frame.DecodePayload(&prepare) != nil || !proto.ValidWorkspaceReadPreparation(prepare.Configuration) || prepare.Configuration.LocalEnvironment == nil || prepare.Configuration.LocalEnvironment.ID != environment.ID { t.Fatal("artifact preparation lost exact local authority", err) } diff --git a/services/core/internal/store/environment_setup.go b/services/core/internal/store/environment_setup.go index 6ba986f7b..04c2e0c73 100644 --- a/services/core/internal/store/environment_setup.go +++ b/services/core/internal/store/environment_setup.go @@ -3,15 +3,12 @@ package store import ( "context" "encoding/json" - "errors" "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/environmentconfig" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" ) @@ -23,20 +20,6 @@ func (s *Store) sealEnvironmentSetup(tenant, resource, id, field string, input a return s.credentialCipher.SealEnvironmentSetup(plaintext, credentialcrypto.EnvironmentSetupBinding{TenantID: tenant, Resource: resource, OwnerID: id, Field: field}) } -func (s *Store) openEnvironmentSetup(tenant, resource, id, field string, ciphertext []byte, output any) error { - if len(ciphertext) == 0 { - return nil - } - plaintext, err := s.credentialCipher.OpenEnvironmentSetup(ciphertext, credentialcrypto.EnvironmentSetupBinding{TenantID: tenant, Resource: resource, OwnerID: id, Field: field}) - if err != nil { - return err - } - if environmentconfig.Decode(plaintext, output) != nil { - return sessions.ErrInvalidInput - } - return nil -} - func (s *Store) saveEnvironmentSetup(ctx context.Context, q *sqlc.Queries, tenant string, session pgtype.UUID, setup environmentconfig.Setup) error { if setup.ValidateInstalled() != nil { return sessions.ErrInvalidInput @@ -54,25 +37,3 @@ func (s *Store) saveEnvironmentSetup(ctx context.Context, q *sqlc.Queries, tenan } return q.CreateEnvironmentSetup(ctx, sqlc.CreateEnvironmentSetupParams{SessionID: session, Contents: encrypted}) } - -func (s *Store) ReadEnvironmentSetup(ctx context.Context, tenant, session string) (environmentconfig.Setup, error) { - var result environmentconfig.Setup - lookup, err := sessionpg.DeviceLookup(tenant, session) - if err != nil { - return result, sessions.ErrNotFound - } - encrypted, err := s.queries.GetEnvironmentSetup(ctx, sqlc.GetEnvironmentSetupParams{TenantID: lookup.TenantID, ID: lookup.ID}) - if errors.Is(err, pgx.ErrNoRows) { - return result, sessions.ErrNotFound - } - if err != nil { - return result, err - } - if err = s.openEnvironmentSetup(uuid.UUID(lookup.TenantID.Bytes).String(), "session", uuid.UUID(lookup.ID.Bytes).String(), "initialization", encrypted, &result); err != nil { - return result, err - } - if result.ValidateInstalled() != nil { - return result, sessions.ErrInvalidInput - } - return result, nil -} diff --git a/services/core/internal/store/environment_setup_test.go b/services/core/internal/store/environment_setup_test.go index 00de50639..a7e9b4134 100644 --- a/services/core/internal/store/environment_setup_test.go +++ b/services/core/internal/store/environment_setup_test.go @@ -31,11 +31,11 @@ func TestEnvironmentSetupEncryptedSnapshotAndIsolation(t *testing.T) { if bytes.Contains(session.Configuration, []byte("canary")) { t.Fatal("plaintext Session metadata") } - frozen, err := s.ReadEnvironmentSetup(t.Context(), tenant, session.ID) + frozen, err := sessionAdapter(s).ReadEnvironmentSetup(t.Context(), tenant, session.ID) if err != nil || !reflect.DeepEqual(frozen, setup) { t.Fatal("Session did not freeze setup", err) } - if _, err := s.ReadEnvironmentSetup(t.Context(), foreign, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).ReadEnvironmentSetup(t.Context(), foreign, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign Session initialization", err) } retry, err := s.CreateSession(t.Context(), tenant, input) diff --git a/services/core/internal/store/environment_skill_references_test.go b/services/core/internal/store/environment_skill_references_test.go index e575e51de..d57c021fd 100644 --- a/services/core/internal/store/environment_skill_references_test.go +++ b/services/core/internal/store/environment_skill_references_test.go @@ -65,7 +65,7 @@ func TestSkillReferencesFreezeWithinSessionCreation(t *testing.T) { } assertFrozen := func(id, version string, archive []byte) { t.Helper() - setup, err := s.ReadEnvironmentSetup(t.Context(), tenant, id) + setup, err := sessionAdapter(s).ReadEnvironmentSetup(t.Context(), tenant, id) if err != nil || len(setup.Skills) != 1 || setup.Skills[0].Metadata.Type != "skill_reference" || setup.Skills[0].Metadata.SkillID != skill.ID || setup.Skills[0].Metadata.Version != version || !bytes.Equal(setup.Skills[0].Archive, archive) { t.Fatal("incorrect frozen installation", err) } @@ -119,7 +119,7 @@ func TestSkillReferencesFreezeWithinSessionCreation(t *testing.T) { t.Fatal("committed retry read deleted sources", err) } assertFrozen(sessionID, "1", first) - if _, err = s.ReadEnvironmentSetup(t.Context(), uuid.NewString(), sessionID); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).ReadEnvironmentSetup(t.Context(), uuid.NewString(), sessionID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign tenant read frozen Skill", err) } } diff --git a/services/core/internal/store/environment_skills_test.go b/services/core/internal/store/environment_skills_test.go index 0172bc20e..52a689a0c 100644 --- a/services/core/internal/store/environment_skills_test.go +++ b/services/core/internal/store/environment_skills_test.go @@ -40,11 +40,11 @@ func TestSkillsFrozenInSession(t *testing.T) { if err != nil { t.Fatal(err) } - frozen, err := s.ReadEnvironmentSetup(t.Context(), tenant, session.ID) + frozen, err := sessionAdapter(s).ReadEnvironmentSetup(t.Context(), tenant, session.ID) if err != nil || !reflect.DeepEqual(frozen.Skills, setup.Skills) { t.Fatal("frozen content changed", err) } - if _, err = s.ReadEnvironmentSetup(t.Context(), foreign, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).ReadEnvironmentSetup(t.Context(), foreign, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign content", err) } } diff --git a/services/core/internal/store/environment_steering_order_test.go b/services/core/internal/store/environment_steering_order_test.go index 32f895500..cff36bdbe 100644 --- a/services/core/internal/store/environment_steering_order_test.go +++ b/services/core/internal/store/environment_steering_order_test.go @@ -3,9 +3,10 @@ package store import ( "context" "errors" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "testing" "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) func TestEnvironmentActiveInputSerializesWithCompletion(t *testing.T) { diff --git a/services/core/internal/store/environment_work_test.go b/services/core/internal/store/environment_work_test.go index 4e85ec56c..2ba9a501c 100644 --- a/services/core/internal/store/environment_work_test.go +++ b/services/core/internal/store/environment_work_test.go @@ -41,7 +41,7 @@ func TestEnvironmentInputWorkFiltersAndPagesDevices(t *testing.T) { default: runtime := h.environments[pending.SessionID] if state == "revoked" { - if err := h.s.RevokeDevice(t.Context(), h.tenant, runtime.device.ID); err != nil { + if err := fixtureSessionService(t, h.db).RevokeDevice(t.Context(), h.tenant, runtime.device.ID); err != nil { t.Fatal(err) } work, err := h.s.ListEnvironmentInputWork(t.Context(), "", []string{runtime.device.ID}) diff --git a/services/core/internal/store/environments.go b/services/core/internal/store/environments.go deleted file mode 100644 index 5907788ac..000000000 --- a/services/core/internal/store/environments.go +++ /dev/null @@ -1,60 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "fmt" - - "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/persistence/postgres/pgunit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -func createSessionEnvironment(ctx context.Context, q *sqlc.Queries, session sqlc.Session) error { - var snapshot struct { - Environment *struct { - Type string `json:"type"` - } `json:"environment"` - } - if err := json.Unmarshal(session.Configuration, &snapshot); err != nil { - return fmt.Errorf("%w: invalid environment configuration", sessions.ErrInvalidInput) - } - if snapshot.Environment == nil || snapshot.Environment.Type == "none" { - return nil - } - switch snapshot.Environment.Type { - case "self_hosted", "openai_hosted": - return q.CreateEnvironment(ctx, sqlc.CreateEnvironmentParams{ - ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, SessionID: session.ID, - }) - default: - return fmt.Errorf("%w: unsupported environment type", sessions.ErrInvalidInput) - } -} - -func (s *Store) GetEnvironment(ctx context.Context, tenantID, environmentID string) (sessions.Environment, error) { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.Environment{}, err - } - id := pgunit.PathID(environmentID) - row, err := s.queries.GetEnvironment(ctx, sqlc.GetEnvironmentParams{TenantID: tenant, ID: id}) - return sessionpg.EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, err) -} - -func (s *Store) GetSessionEnvironment(ctx context.Context, tenantID, sessionID string) (sessions.Environment, error) { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.Environment{}, err - } - id, err := parseID(sessionID) - if err != nil { - return sessions.Environment{}, err - } - row, err := s.queries.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: tenant, ID: id}) - return sessionpg.EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, err) -} diff --git a/services/core/internal/store/environments_test.go b/services/core/internal/store/environments_test.go index b82284250..d423b4156 100644 --- a/services/core/internal/store/environments_test.go +++ b/services/core/internal/store/environments_test.go @@ -34,17 +34,17 @@ func TestEnvironmentOwnershipPersistsAndStaysScoped(t *testing.T) { if err != nil { t.Fatal(err) } - first, err := s.GetSessionEnvironment(ctx, tenant, session.ID) + first, err := sessionAdapter(s).GetSessionEnvironment(ctx, tenant, session.ID) if err != nil || first.ID == "" || first.ID == session.ID || first.SessionID != session.ID || first.TenantID != tenant || first.Status != "pending" || first.CreatedAt.IsZero() { t.Fatal(first, err) } - got, err := s.GetEnvironment(ctx, tenant, first.ID) + got, err := sessionAdapter(s).GetEnvironment(ctx, tenant, first.ID) if err != nil || !reflect.DeepEqual(got, first) { t.Fatal(got, err) } for _, lookup := range []func() error{ - func() error { _, err := s.GetEnvironment(ctx, foreign, first.ID); return err }, - func() error { _, err := s.GetSessionEnvironment(ctx, foreign, session.ID); return err }, + func() error { _, err := sessionAdapter(s).GetEnvironment(ctx, foreign, first.ID); return err }, + func() error { _, err := sessionAdapter(s).GetSessionEnvironment(ctx, foreign, session.ID); return err }, } { if err := lookup(); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign access", err) @@ -54,7 +54,7 @@ func TestEnvironmentOwnershipPersistsAndStaysScoped(t *testing.T) { if err != nil { t.Fatal(err) } - otherEnvironment, err := s.GetSessionEnvironment(ctx, foreign, other.ID) + otherEnvironment, err := sessionAdapter(s).GetSessionEnvironment(ctx, foreign, other.ID) if err != nil || otherEnvironment.ID == first.ID { t.Fatal(otherEnvironment, err) } @@ -71,7 +71,7 @@ func TestEnvironmentOwnershipPersistsAndStaysScoped(t *testing.T) { if err != nil || retry.ID != session.ID || retry.Metadata["updated"] != "yes" { t.Fatal(retry, err) } - retained, err := restarted.GetSessionEnvironment(ctx, tenant, retry.ID) + retained, err := sessionAdapter(restarted).GetSessionEnvironment(ctx, tenant, retry.ID) if err != nil || !reflect.DeepEqual(retained, first) { t.Fatal(retained, err) } @@ -108,7 +108,7 @@ func TestEnvironmentCreationWinnerOwnsSnapshotAndIdentity(t *testing.T) { t.Error(err) return } - environment, err := st.GetSessionEnvironment(ctx, tenant, creation.Session.ID) + environment, err := sessionAdapter(st).GetSessionEnvironment(ctx, tenant, creation.Session.ID) if err != nil { t.Error(err) return @@ -170,7 +170,7 @@ func TestEnvironmentCreationWinnerOwnsSnapshotAndIdentity(t *testing.T) { if err != nil || found.Created || found.Session.ID != first.creation.Session.ID { t.Fatal(found, err) } - environment, err := restarted.GetSessionEnvironment(ctx, tenant, found.Session.ID) + environment, err := sessionAdapter(restarted).GetSessionEnvironment(ctx, tenant, found.Session.ID) if err != nil || !reflect.DeepEqual(environment, first.environment) { t.Fatal(environment, err) } @@ -222,7 +222,7 @@ func TestEnvironmentCreationFailureRollsBackAllResources(t *testing.T) { t.Fatal("retry remained reserved", session, err) } environmentInputHistory(t, pool, session.ID, 0, 0) - if environment, err := s.GetSessionEnvironment(ctx, tenant, session.ID); err != nil || environment.ID == "" { + if environment, err := sessionAdapter(s).GetSessionEnvironment(ctx, tenant, session.ID); err != nil || environment.ID == "" { t.Fatal(environment, err) } }) @@ -238,17 +238,17 @@ func TestEnvironmentDeletionHidesWithoutDestroyingOwnership(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(ctx, tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } if err := s.DeleteSession(ctx, tenant, session.ID); err != nil { t.Fatal(err) } - if _, err := s.GetEnvironment(ctx, tenant, environment.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetEnvironment(ctx, tenant, environment.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - if _, err := s.GetSessionEnvironment(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSessionEnvironment(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } if _, err := s.CreateSession(ctx, tenant, input); !errors.Is(err, sessions.ErrIdempotencyConflict) { @@ -270,7 +270,7 @@ func TestEnvironmentAbsentForNoneAndLegacySnapshots(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.GetSessionEnvironment(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSessionEnvironment(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unexpected Environment", err) } retry, err := s.CreateSession(ctx, tenant, input) diff --git a/services/core/internal/store/execution_test.go b/services/core/internal/store/execution_test.go index 2c8e5a92e..1ccecaa1e 100644 --- a/services/core/internal/store/execution_test.go +++ b/services/core/internal/store/execution_test.go @@ -76,11 +76,11 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { tenant, active := newTurnSession(t, s) input := submitMessage(t, s, tenant, active.ID, "active") transition(t, writer, tenant, active.ID, input.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - host, err := s.CreateDevice(t.Context(), tenant, "owner test", runtimedevice.HashCredential(uuid.NewString())) + host, err := sessionService(t, s).CreateDevice(t.Context(), tenant, "owner test", runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - if err = writer.BindSessionDevice(t.Context(), tenant, active.ID, host.ID); err != nil { + if err = sessionExecution(t, writer.lease).BindSessionDevice(t.Context(), tenant, active.ID, host.ID); err != nil { t.Fatal(err) } queued, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "queued"}) @@ -121,7 +121,7 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { t.Fatalf("stale %s committed while successor owned lease", name) } } - mustReject("binding", writer.BindSessionDevice(t.Context(), tenant, queued.ID, host.ID)) + mustReject("binding", sessionExecution(t, writer.lease).BindSessionDevice(t.Context(), tenant, queued.ID, host.ID)) _, err = writer.TransitionTurn(t.Context(), tenant, queued.ID, pending.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) mustReject("claim", err) mustReject("journal", operations.AppendTurnEvents(t.Context(), tenant, active.ID, input.TurnID, 1, []sessions.ExecutionEvent{{Kind: "delta", Payload: json.RawMessage(`{"delta":"stale"}`)}})) @@ -150,7 +150,7 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { if err != nil || saved.Applied { t.Fatal("stale receipt persisted", saved, err) } - if _, err = s.GetSessionDevice(t.Context(), tenant, queued.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).GetSessionDevice(t.Context(), tenant, queued.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("stale binding persisted", err) } queuedTurn, err := s.GetTurn(t.Context(), tenant, queued.ID, pending.TurnID) @@ -176,7 +176,7 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { if err = successor.lease.Close(t.Context()); err != nil { t.Fatal(err) } - mustReject("closed writer", successor.BindSessionDevice(t.Context(), tenant, queued.ID, host.ID)) + mustReject("closed writer", sessionExecution(t, successor.lease).BindSessionDevice(t.Context(), tenant, queued.ID, host.ID)) } func TestExecutionWriterSerializesWritesOnItsLease(t *testing.T) { @@ -262,11 +262,9 @@ 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, - "reconciliation": s.ReconcileEnvironmentConnections(t.Context()), - "connection change": s.ReplaceEnvironmentConnection(t.Context(), tenant, uuid.NewString(), uuid.NewString()), + "ownership check": s.checkExecutionOwnership(t.Context()), + "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/executor_credential_target.go b/services/core/internal/store/executor_credential_target.go index b6930966d..3a3796866 100644 --- a/services/core/internal/store/executor_credential_target.go +++ b/services/core/internal/store/executor_credential_target.go @@ -10,6 +10,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -50,7 +51,7 @@ func (s *Store) withExecutorCredentialTarget(ctx context.Context, principal iden if !environment.Valid { return s.pooled.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { return apply(ctx, s.queries.WithTx(tx)) }) } - owned, err := s.GetEnvironment(ctx, principal.TenantID, uuid.UUID(environment.Bytes).String()) + owned, err := sessionpg.LoadEnvironment(ctx, s.queries, principal.TenantID, uuid.UUID(environment.Bytes).String()) if err != nil { return err } diff --git a/services/core/internal/store/executor_principals_migration_test.go b/services/core/internal/store/executor_principals_migration_test.go index 83b06d872..b2d5bc876 100644 --- a/services/core/internal/store/executor_principals_migration_test.go +++ b/services/core/internal/store/executor_principals_migration_test.go @@ -105,7 +105,7 @@ func TestExecutorPrincipalMigrationRetiresUnknownAuthority(t *testing.T) { t.Cleanup(migrated.Close) s := New(migrated) p := FixtureExecutorPrincipal(t, s, tenant) - if _, err := s.AuthenticateEnvironmentExecutor(ctx, environment, digest); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, environment, digest); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("legacy credential accepted", err) } if _, err := s.IssueExecutorCredential(ctx, p, environment, ""); !errors.Is(err, sessions.ErrExecutorCredentialExists) { diff --git a/services/core/internal/store/executor_principals_test.go b/services/core/internal/store/executor_principals_test.go index 9712a5936..3b9ecc4b8 100644 --- a/services/core/internal/store/executor_principals_test.go +++ b/services/core/internal/store/executor_principals_test.go @@ -30,7 +30,7 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(ctx, principal.TenantID, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(ctx, principal.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -38,7 +38,7 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.T) { } check := func(st *Store, environment string, key sessions.IssuedExecutorCredential, allowed bool) { t.Helper() - owner, err := st.AuthenticateEnvironmentExecutor(ctx, environment, executorDigest(key.Token)) + owner, err := sessionAdapter(st).AuthenticateEnvironmentExecutor(ctx, environment, executorDigest(key.Token)) if allowed && (err != nil || owner != p.TenantID) || !allowed && !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unexpected principal authorization", allowed, err) } @@ -145,7 +145,7 @@ func TestExecutorPrincipalRequiresVerifiedScopeAndRecordedCreator(t *testing.T) if err != nil { t.Fatal(err) } - target, err := s.GetSessionEnvironment(ctx, p.TenantID, session.ID) + target, err := sessionAdapter(s).GetSessionEnvironment(ctx, p.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -153,13 +153,13 @@ func TestExecutorPrincipalRequiresVerifiedScopeAndRecordedCreator(t *testing.T) if err != nil { t.Fatal(err) } - if owner, err := s.AuthenticateEnvironmentExecutor(ctx, target.ID, executorDigest(key.Token)); err != nil || owner != p.TenantID { + if owner, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, target.ID, executorDigest(key.Token)); err != nil || owner != p.TenantID { t.Fatal("user principal failed", err) } if _, err := pool.Exec(ctx, "UPDATE sessions SET creator_kind=NULL,creator_id=NULL WHERE id=$1", session.ID); err != nil { t.Fatal(err) } - if _, err := s.AuthenticateEnvironmentExecutor(ctx, target.ID, executorDigest(key.Token)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).AuthenticateEnvironmentExecutor(ctx, target.ID, executorDigest(key.Token)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unknown creator accepted", err) } if _, err := s.IssueExecutorCredential(ctx, p, uuid.NewString(), target.ID); !errors.Is(err, sessions.ErrNotFound) { diff --git a/services/core/internal/store/function_execution_native_test.go b/services/core/internal/store/function_execution_native_test.go index 99a60a802..58458ac01 100644 --- a/services/core/internal/store/function_execution_native_test.go +++ b/services/core/internal/store/function_execution_native_test.go @@ -19,7 +19,7 @@ func TestNativeFunctionExecutionPersistsCallsResultsAndContinuity(t *testing.T) if err != nil { t.Fatal(err) } - if err := h.s.BindSessionDevice(ctx, h.tenant, h.session.ID, h.device.ID); err != nil { + if err := bindSessionDevice(t, h.db, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } nativeID := "" diff --git a/services/core/internal/store/function_execution_test.go b/services/core/internal/store/function_execution_test.go index 58aa56681..b35aec657 100644 --- a/services/core/internal/store/function_execution_test.go +++ b/services/core/internal/store/function_execution_test.go @@ -25,7 +25,7 @@ func newFunctionHarness(t *testing.T) *dispatchHarness { if err != nil { t.Fatal(err) } - if err := h.s.BindSessionDevice(t.Context(), h.tenant, h.session.ID, h.device.ID); err != nil { + if err := bindSessionDevice(t, h.db, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, FunctionTools: proto.CapabilitySupported, FunctionResultImages: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) @@ -195,7 +195,7 @@ func TestExecutionFunctionsRequireAdvertisedCapability(t *testing.T) { t.Fatal(err) } h.session = session - if err := h.s.BindSessionDevice(t.Context(), h.tenant, session.ID, h.device.ID); err != nil { + if err := bindSessionDevice(t, h.db, h.tenant, session.ID, h.device.ID); err != nil { t.Fatal(err) } input := h.message("start", "Run") diff --git a/services/core/internal/store/function_worker_test.go b/services/core/internal/store/function_worker_test.go index 4f630b84c..00cfb8382 100644 --- a/services/core/internal/store/function_worker_test.go +++ b/services/core/internal/store/function_worker_test.go @@ -43,7 +43,7 @@ func TestWorkerWaitsForToolCapabilities(t *testing.T) { } } if prebound && isMCP { - if err := h.s.BindSessionDevice(t.Context(), h.tenant, h.session.ID, h.device.ID); err != nil { + if err := bindSessionDevice(t, h.db, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } } @@ -84,7 +84,7 @@ func TestWorkerWaitsForToolCapabilities(t *testing.T) { t.Fatal(current, err) } if !prebound { - if _, err := h.s.GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := fixtureSessionStore(h.db).GetSessionDevice(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("bound an incapable device", err) } } diff --git a/services/core/internal/store/hosted_initialization_failure_public_test.go b/services/core/internal/store/hosted_initialization_failure_public_test.go index c34a1bbb2..6ea80d866 100644 --- a/services/core/internal/store/hosted_initialization_failure_public_test.go +++ b/services/core/internal/store/hosted_initialization_failure_public_test.go @@ -109,7 +109,7 @@ func hostedFailureStore(t *testing.T) (*store.Store, fixtureDB) { return store.NewWithCredentialCipher(pool, cipher), fixtureDB{pool: pool, cipher: cipher} } -func hostedFailureSession(t *testing.T, s *store.Store, tenant string, input sessions.CreateSession) (sessions.Session, sessions.Environment) { +func hostedFailureSession(t *testing.T, s *store.Store, db fixtureDB, tenant string, input sessions.CreateSession) (sessions.Session, sessions.Environment) { t.Helper() input.Creator, input.Engine, input.IdempotencyKey = store.FixtureCreator(), "codex", uuid.NewString() input.Configuration = json.RawMessage(`{"agent":{"id":"agent_test","model":"test-model","tools":[]},"environment":{"type":"openai_hosted","network":{"access":"enabled"}}}`) @@ -120,7 +120,7 @@ func hostedFailureSession(t *testing.T, s *store.Store, tenant string, input ses if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -134,7 +134,7 @@ func failHostedInitialization(t *testing.T, s *store.Store, db fixtureDB, tenant if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err != nil { t.Fatal(err) } - awaitInitialization(t, s, tenant, environment.ID, "failed") + awaitInitialization(t, db.pool, tenant, environment.ID, "failed") } // H1/H2/H3/H4: one transaction records the Environment failure, an error event @@ -183,7 +183,7 @@ func TestHostedInitializationFailureRecordsSafeSessionFailure(t *testing.T) { t.Run(test.name, func(t *testing.T) { s, db := hostedFailureStore(t) tenant := uuid.NewString() - session, environment := hostedFailureSession(t, s, tenant, test.input) + session, environment := hostedFailureSession(t, s, db, tenant, test.input) p := &hostedFailureProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, fail: test.p.fail, skip: test.p.skip, result: test.p.result, err: test.p.err} failHostedInitialization(t, s, db, tenant, environment, p) @@ -248,7 +248,7 @@ func TestHostedInitializationFailureRecordsSafeSessionFailure(t *testing.T) { func TestHostedInitializationFailureSettlesPendingInitialInput(t *testing.T) { s, db := hostedFailureStore(t) tenant := uuid.NewString() - session, environment := hostedFailureSession(t, s, tenant, sessions.CreateSession{ + session, environment := hostedFailureSession(t, s, db, tenant, sessions.CreateSession{ Initialization: environmentconfig.Setup{Commands: []environmentconfig.SetupCommand{{Command: "exit 3"}}}, InitialInputs: []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"initial"}`)}}, }) @@ -282,7 +282,7 @@ func TestHostedInitializationFailureSettlesPendingInitialInput(t *testing.T) { func TestHostedInitializationFailurePublicHTTP(t *testing.T) { s, db := hostedFailureStore(t) tenant, token, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() - session, environment := hostedFailureSession(t, s, tenant, sessions.CreateSession{ + session, environment := hostedFailureSession(t, s, db, tenant, sessions.CreateSession{ Initialization: environmentconfig.Setup{Commands: []environmentconfig.SetupCommand{{Command: "echo " + hostedFailureCanary + "; exit 3"}}}, Metadata: map[string]string{"case": "setup-exit3"}, }) @@ -336,7 +336,7 @@ func TestHostedInitializationFailurePublicHTTP(t *testing.T) { if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err != nil { t.Fatal(err) } - awaitInitialization(t, s, tenant, environment.ID, "failed") + awaitInitialization(t, db.pool, tenant, environment.ID, "failed") reason := `Failed to provision environment: script "setup_commands[0]" failed with exit code 3` read, err := s.GetSession(t.Context(), tenant, session.ID) diff --git a/services/core/internal/store/initial_files.go b/services/core/internal/store/initial_files.go index 9b8e7f2da..6956d6fb4 100644 --- a/services/core/internal/store/initial_files.go +++ b/services/core/internal/store/initial_files.go @@ -10,7 +10,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/files" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -82,21 +81,3 @@ func readInitialSourceFile(ctx context.Context, tx pgx.Tx, source sqlc.SourceFil } return body, reader.Close() } - -// ReadInitialEnvironmentFile decrypts only the next frozen file, bounding memory per installation. -func (s *Store) ReadInitialEnvironmentFile(ctx context.Context, tenant, session string, position int) (environmentconfig.InitialFileMetadata, []byte, error) { - lookup, err := sessionpg.DeviceLookup(tenant, session) - if err != nil { - return environmentconfig.InitialFileMetadata{}, nil, err - } - row, err := s.queries.GetInitialEnvironmentFile(ctx, sqlc.GetInitialEnvironmentFileParams{TenantID: lookup.TenantID, SessionID: lookup.ID, Position: int32(position)}) - if err != nil { - return environmentconfig.InitialFileMetadata{}, nil, err - } - id := uuid.UUID(row.ID.Bytes).String() - body, err := s.credentialCipher.OpenEnvironmentFile(row.Contents, credentialcrypto.EnvironmentFileBinding{TenantID: uuid.UUID(lookup.TenantID.Bytes).String(), Resource: "session", OwnerID: uuid.UUID(lookup.ID.Bytes).String(), FileID: id}) - if err == nil && int64(len(body)) != row.SizeBytes { - err = sessions.ErrInvalidInput - } - return environmentconfig.InitialFileMetadata{ID: id, Path: row.Path, SizeBytes: &row.SizeBytes}, body, err -} diff --git a/services/core/internal/store/initial_files_http_test.go b/services/core/internal/store/initial_files_http_test.go index 4759a6ae3..beb136f88 100644 --- a/services/core/internal/store/initial_files_http_test.go +++ b/services/core/internal/store/initial_files_http_test.go @@ -59,7 +59,7 @@ func TestInitialFilesHTTPInlineLimitsAndRetry(t *testing.T) { id = response.ID } for position := range files { - _, actual, err := s.ReadInitialEnvironmentFile(t.Context(), tenant, id, position) + _, actual, err := fixtureSessionStore(db).ReadInitialEnvironmentFile(t.Context(), tenant, id, position) if err != nil || !bytes.Equal(actual, data) { t.Fatal("large HTTP snapshot differs", err) } diff --git a/services/core/internal/store/initial_files_test.go b/services/core/internal/store/initial_files_test.go index fbc56225c..aaf637331 100644 --- a/services/core/internal/store/initial_files_test.go +++ b/services/core/internal/store/initial_files_test.go @@ -67,11 +67,11 @@ func TestInitialFilesFrozenEncryptedIsolatedAndRetryable(t *testing.T) { t.Fatal("retry re-resolved deleted resources", err) } for position := range initial { - metadata, body, err := s.ReadInitialEnvironmentFile(t.Context(), strings.ToUpper(tenant), strings.ToUpper(session.ID), position) + metadata, body, err := sessionAdapter(s).ReadInitialEnvironmentFile(t.Context(), strings.ToUpper(tenant), strings.ToUpper(session.ID), position) if err != nil || !bytes.Equal(body, canary) || metadata.ID == "" { t.Fatal("frozen initial content", err) } - if _, _, err := s.ReadInitialEnvironmentFile(t.Context(), foreign, session.ID, position); err == nil { + if _, _, err := sessionAdapter(s).ReadInitialEnvironmentFile(t.Context(), foreign, session.ID, position); err == nil { t.Fatal("foreign bytes disclosed") } var encrypted []byte @@ -85,7 +85,7 @@ func TestInitialFilesFrozenEncryptedIsolatedAndRetryable(t *testing.T) { if _, err := s.CreateSession(t.Context(), tenant, changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal("changed bytes retried", err) } - if _, err := s.GetSessionDevice(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSessionDevice(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("uninitialized environment exposed", err) } } diff --git a/services/core/internal/store/list_cursor_public_test.go b/services/core/internal/store/list_cursor_public_test.go index 2cafa6538..8ad3e30bd 100644 --- a/services/core/internal/store/list_cursor_public_test.go +++ b/services/core/internal/store/list_cursor_public_test.go @@ -14,6 +14,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/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/skills" @@ -34,7 +35,7 @@ type cursorFixture struct { file string } -func seedCursorFixture(t *testing.T, s *store.Store, writer *store.Store, journal *sessions.ExecutionOperations, skillService *skills.Service, artifacts *sessions.Service, client pathIDClient, token, tenant, label string) cursorFixture { +func seedCursorFixture(t *testing.T, s *store.Store, leased execution.Owner, skillService *skills.Service, sessionService *sessions.Service, client pathIDClient, token, tenant, label string) cursorFixture { t.Helper() ctx := t.Context() var f cursorFixture @@ -76,10 +77,10 @@ func seedCursorFixture(t *testing.T, s *store.Store, writer *store.Store, journa artifactSession, environment := hostedArtifactSession(t, s, tenant, label+"-artifacts") f.artifactSession = artifactSession - f.artifactTurn = completeArtifactTurn(t, s, artifacts, tenant, artifactSession, environment, label+"-artifact-turn", map[string]string{"a.txt": "alpha", "c.txt": "charlie"}) + f.artifactTurn = completeArtifactTurn(t, s, sessionService, tenant, artifactSession, environment, label+"-artifact-turn", map[string]string{"a.txt": "alpha", "c.txt": "charlie"}) f.artifact = first("/v1/agents/sessions/" + artifactSession + "/artifacts") otherArtifactSession, otherEnvironment := hostedArtifactSession(t, s, tenant, label+"-other-artifacts") - completeArtifactTurn(t, s, artifacts, tenant, otherArtifactSession, otherEnvironment, label+"-other-artifact-turn", map[string]string{"b.txt": "bravo"}) + completeArtifactTurn(t, s, sessionService, tenant, otherArtifactSession, otherEnvironment, label+"-other-artifact-turn", map[string]string{"b.txt": "bravo"}) f.otherArtifact = first("/v1/agents/sessions/" + otherArtifactSession + "/artifacts") // Subagent history is seeded through the execution lease, as a daemon would. @@ -94,14 +95,14 @@ func seedCursorFixture(t *testing.T, s *store.Store, writer *store.Store, journa if err != nil { t.Fatal(err) } - host, err := s.CreateDevice(ctx, tenant, "cursor "+key, runtimedevice.HashCredential(uuid.NewString())) + host, err := sessionService.CreateDevice(ctx, tenant, "cursor "+key, runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - if err = writer.BindSessionDevice(ctx, tenant, created.ID, host.ID); err != nil { + if err = leased.Sessions.BindSessionDevice(ctx, tenant, created.ID, host.ID); err != nil { t.Fatal(err) } - if _, err = writer.TransitionTurn(ctx, tenant, created.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err = leased.Store.TransitionTurn(ctx, tenant, created.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } opened := int64(1700000001000) @@ -121,7 +122,7 @@ func seedCursorFixture(t *testing.T, s *store.Store, writer *store.Store, journa turn("child", "child-turn", opened), message("child", "child-turn", "child-item", 0), message("child", "child-turn", "child-item-2", 1), turn("child", "later-child-turn", opened+1000), message("child", "later-child-turn", "later-child-item", 0), turn("sibling", "sibling-turn", opened), message("sibling", "sibling-turn", "sibling-item", 0)} - if err = journal.AppendTurnEvents(ctx, tenant, created.ID, receipt.TurnID, 1, facts); err != nil { + if err = leased.Sessions.AppendTurnEvents(ctx, tenant, created.ID, receipt.TurnID, 1, facts); err != nil { t.Fatal(err) } return created.ID, receipt.TurnID @@ -240,14 +241,13 @@ func TestListCursorErrorsPostgres(t *testing.T) { defer server.Close() client := pathIDClient{t: t, server: server} leased := executionOwner(t, db, s) - writer := leased.Store skillService := store.SkillService(t, db.pool, db.cipher) _, sessionService, err := fixtureSessions(db) if err != nil { t.Fatal(err) } - a := seedCursorFixture(t, s, writer, leased.Sessions, skillService, sessionService, client, owner, ownerTenant, "a") - b := seedCursorFixture(t, s, writer, leased.Sessions, skillService, sessionService, client, foreign, foreignTenant, "b") + a := seedCursorFixture(t, s, leased, skillService, sessionService, client, owner, ownerTenant, "a") + b := seedCursorFixture(t, s, leased, skillService, sessionService, client, foreign, foreignTenant, "b") text := func(value string) *string { return &value } var ( diff --git a/services/core/internal/store/local_environment_devices.go b/services/core/internal/store/local_environment_devices.go deleted file mode 100644 index 684dbfeec..000000000 --- a/services/core/internal/store/local_environment_devices.go +++ /dev/null @@ -1,38 +0,0 @@ -package store - -import ( - "context" - - "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/google/uuid" - "github.com/jackc/pgx/v5/pgtype" -) - -// CreateEnvironmentDevice provisions one dedicated Runtime without widening an existing credential. -func (s *Store) CreateEnvironmentDevice(ctx context.Context, tenantID, environmentID, name, credentialHash string) (sessions.ExecutionDevice, error) { - environment, err := s.GetEnvironment(ctx, tenantID, environmentID) - if err != nil { - return sessions.ExecutionDevice{}, err - } - if kind, err := sessions.EnvironmentType(environment.Configuration); err != nil || kind != "openai_hosted" { - return sessions.ExecutionDevice{}, sessions.ErrInvalidInput - } - lookup, err := sessionpg.DeviceLookup(tenantID, environment.ID) - if err != nil { - return sessions.ExecutionDevice{}, err - } - params, err := newDeviceParams(lookup.TenantID, name, credentialHash) - if err != nil { - return sessions.ExecutionDevice{}, err - } - device := sessions.ExecutionDevice{ID: uuid.UUID(params.ID.Bytes).String(), Name: params.Name, EnvironmentID: environment.ID} - err = s.withPublicSession(ctx, tenantID, environment.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - return sessions.CreateEnvironmentDevice(ctx, sessionpg.BindSession(q, lookup.TenantID, session), device, params.CredentialHash.String) - }) - if err != nil { - return sessions.ExecutionDevice{}, err - } - return device, nil -} diff --git a/services/core/internal/store/local_environment_devices_test.go b/services/core/internal/store/local_environment_devices_test.go index 25ec16580..00c31eff8 100644 --- a/services/core/internal/store/local_environment_devices_test.go +++ b/services/core/internal/store/local_environment_devices_test.go @@ -20,7 +20,7 @@ func localEnvironment(t *testing.T, s *Store, tenant string) (sessions.Session, if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -28,21 +28,22 @@ func localEnvironment(t *testing.T, s *Store, tenant string) (sessions.Session, } func TestEnvironmentDeviceAuthorityAndLifecycle(t *testing.T) { - s, _ := testStore(t) + s, pool := testStore(t) tenant, foreignTenant := uuid.NewString(), uuid.NewString() session, environment := localEnvironment(t, s, tenant) sibling, _ := localEnvironment(t, s, tenant) foreign, _ := localEnvironment(t, s, foreignTenant) digest := runtimedevice.HashCredential(uuid.NewString()) - if _, err := s.CreateEnvironmentDevice(t.Context(), foreignTenant, environment.ID, "foreign", digest); !errors.Is(err, sessions.ErrNotFound) { + if _, err := FixtureEnvironmentDevice(t.Context(), pool, foreignTenant, environment.ID, "foreign", digest); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("foreign provisioning: %v", err) } - bound, err := s.CreateEnvironmentDevice(t.Context(), tenant, environment.ID, "dedicated", digest) + bound, err := FixtureEnvironmentDevice(t.Context(), pool, tenant, environment.ID, "dedicated", digest) if err != nil || bound.EnvironmentID != environment.ID { t.Fatalf("provision: %+v %v", bound, err) } + execution := sessionExecution(t, executionWriter(t, s).lease) for _, other := range []sessions.Session{sibling, foreign} { - if err := s.BindSessionDevice(t.Context(), other.TenantID, other.ID, bound.ID); err == nil { + if err := execution.BindSessionDevice(t.Context(), other.TenantID, other.ID, bound.ID); err == nil { t.Fatal("dedicated credential bound to another Session") } } @@ -51,27 +52,27 @@ func TestEnvironmentDeviceAuthorityAndLifecycle(t *testing.T) { t.Fatalf("dedicated device entered general selection: %v %v", devices, err) } reopened, _ := testStore(t) - got, err := reopened.GetSessionDevice(t.Context(), tenant, session.ID) + got, err := sessionAdapter(reopened).GetSessionDevice(t.Context(), tenant, session.ID) if err != nil || got != bound { t.Fatalf("durable exact binding: %+v %v", got, err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), bound.ID); err != nil || !ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), bound.ID); err != nil || !ok { t.Fatalf("valid credential unavailable: %v", err) } if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { t.Fatal(err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), bound.ID); err != nil || ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), bound.ID); err != nil || ok { t.Fatalf("deleted Environment still authenticates: %v", err) } - status, err := s.TouchRuntimeHeartbeat(t.Context(), bound.ID) + status, err := sessionService(t, s).TouchRuntimeHeartbeat(t.Context(), bound.ID) if err != nil || !status.Deleted { t.Fatalf("deleted Environment heartbeat: %+v %v", status, err) } } func TestEnvironmentDeviceProvisioningHasOneWinner(t *testing.T) { - s, _ := testStore(t) + s, pool := testStore(t) tenant := uuid.NewString() session, environment := localEnvironment(t, s, tenant) var wg sync.WaitGroup @@ -80,7 +81,7 @@ func TestEnvironmentDeviceProvisioningHasOneWinner(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - _, err := s.CreateEnvironmentDevice(t.Context(), tenant, environment.ID, "runtime", runtimedevice.HashCredential(uuid.NewString())) + _, err := FixtureEnvironmentDevice(t.Context(), pool, tenant, environment.ID, "runtime", runtimedevice.HashCredential(uuid.NewString())) results <- err }() } @@ -97,14 +98,14 @@ func TestEnvironmentDeviceProvisioningHasOneWinner(t *testing.T) { if winners != 1 { t.Fatalf("provisioned %d devices", winners) } - bound, err := s.GetSessionDevice(t.Context(), tenant, session.ID) + bound, err := sessionAdapter(s).GetSessionDevice(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } - if err := s.RevokeDevice(t.Context(), tenant, bound.ID); err != nil { + if err := sessionService(t, s).RevokeDevice(t.Context(), tenant, bound.ID); err != nil { t.Fatal(err) } - if _, err := s.CreateEnvironmentDevice(t.Context(), tenant, environment.ID, "replacement", runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, sessions.ErrDeviceBindingConflict) { + if _, err := FixtureEnvironmentDevice(t.Context(), pool, tenant, environment.ID, "replacement", runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, sessions.ErrDeviceBindingConflict) { t.Fatalf("silent placement replacement: %v", err) } } diff --git a/services/core/internal/store/local_environment_worker_test.go b/services/core/internal/store/local_environment_worker_test.go index 8a029669c..1aabca710 100644 --- a/services/core/internal/store/local_environment_worker_test.go +++ b/services/core/internal/store/local_environment_worker_test.go @@ -22,7 +22,7 @@ func localWorker(t *testing.T, scoped, execute bool) (*dispatchHarness, *executi _, pool := store.NewTestStore(t) insertWorkerRuntimeAllocation(t, pool, h, "disabled") } - environment, err := h.s.GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) + environment, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) if err != nil { t.Fatal(err) } @@ -96,14 +96,14 @@ func TestLocalEnvironmentWorkerRejectsGeneralDeviceDespiteCapability(t *testing. if err != nil { t.Fatal(err) } - unassigned, err := h.s.GetSessionEnvironment(t.Context(), h.tenant, other.ID) + unassigned, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, other.ID) if err != nil { t.Fatal(err) } if _, err := w.ReadEnvironmentDirectory(t.Context(), unassigned, "reports"); !errors.Is(err, execution.ErrExecutionUnavailable) { t.Fatal("unassigned environment selected general device", err) } - if _, err := h.s.GetSessionDevice(t.Context(), h.tenant, other.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := fixtureSessionStore(h.db).GetSessionDevice(t.Context(), h.tenant, other.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("read persisted an unauthorized placement", err) } } diff --git a/services/core/internal/store/message_input_helpers_test.go b/services/core/internal/store/message_input_helpers_test.go index e4007b401..550d7ef5c 100644 --- a/services/core/internal/store/message_input_helpers_test.go +++ b/services/core/internal/store/message_input_helpers_test.go @@ -1,8 +1,9 @@ package store_test import ( - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "testing" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" ) func inputTextForTest(t *testing.T, input proto.MessageInput) string { diff --git a/services/core/internal/store/native_environment_test.go b/services/core/internal/store/native_environment_test.go index 69f6f3856..e86f74c55 100644 --- a/services/core/internal/store/native_environment_test.go +++ b/services/core/internal/store/native_environment_test.go @@ -103,7 +103,7 @@ func TestNativeNoExecutionEnvironment(t *testing.T) { if err != nil { t.Fatal(err) } - if err = h.s.BindSessionDevice(ctx, h.tenant, h.session.ID, h.device.ID); err != nil { + if err = bindSessionDevice(t, h.db, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } first := h.message("first", "Return an answer.") diff --git a/services/core/internal/store/native_public_execution_test.go b/services/core/internal/store/native_public_execution_test.go index d22a8da9a..496914f2e 100644 --- a/services/core/internal/store/native_public_execution_test.go +++ b/services/core/internal/store/native_public_execution_test.go @@ -23,7 +23,9 @@ func verifyNativePublicExecution(t *testing.T, h *dispatchHarness, parent contex } ctx, cancel := context.WithCancel(parent) defer cancel() - worker := startWorker(t, ctx, h.db, h.d) + // The Turns before this proof ran on the harness's execution Owner, so the + // Worker takes that lease rather than a second one. + worker := startOwnedWorker(t, ctx, h.db, h.d, h.owner()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { diff --git a/services/core/internal/store/native_recovery_test.go b/services/core/internal/store/native_recovery_test.go index 16064b132..f95a1293e 100644 --- a/services/core/internal/store/native_recovery_test.go +++ b/services/core/internal/store/native_recovery_test.go @@ -2,8 +2,9 @@ package store import ( "errors" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) func TestSessionExecutionBindingRetainsStartedExecutionRequirement(t *testing.T) { @@ -11,7 +12,13 @@ func TestSessionExecutionBindingRetainsStartedExecutionRequirement(t *testing.T) tenant, session := newTurnSession(t, s) foreign, _ := newTurnSession(t, s) device, _ := registerTestDevice(t, s, tenant) - if err := s.BindSessionDevice(t.Context(), tenant, session.ID, device.ID); err != nil { + // The bind runs on an execution lease of its own, which closes before the + // pool does. + writer := executionWriter(t, s) + if err := sessionExecution(t, writer.lease).BindSessionDevice(t.Context(), tenant, session.ID, device.ID); err != nil { + t.Fatal(err) + } + if err := writer.lease.Close(t.Context()); err != nil { t.Fatal(err) } assertStarted := func(st *Store, want bool) { diff --git a/services/core/internal/store/no_environment_test.go b/services/core/internal/store/no_environment_test.go index c81e554df..c43cf3fed 100644 --- a/services/core/internal/store/no_environment_test.go +++ b/services/core/internal/store/no_environment_test.go @@ -16,7 +16,7 @@ func TestNoEnvironmentRejectsUnadvertisedDeviceBeforeClaim(t *testing.T) { if err != nil { t.Fatal(err) } - if err = h.s.BindSessionDevice(ctx, h.tenant, h.session.ID, h.device.ID); err != nil { + if err = bindSessionDevice(t, h.db, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } input := h.message("first", "Answer") diff --git a/services/core/internal/store/provider_operations_fixture_test.go b/services/core/internal/store/provider_operations_fixture_test.go index 68634162e..68d0d0b98 100644 --- a/services/core/internal/store/provider_operations_fixture_test.go +++ b/services/core/internal/store/provider_operations_fixture_test.go @@ -2,6 +2,7 @@ package store_test import ( "context" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/providercontract" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" diff --git a/services/core/internal/store/public_execution_test.go b/services/core/internal/store/public_execution_test.go index 855d52845..9c8df0da2 100644 --- a/services/core/internal/store/public_execution_test.go +++ b/services/core/internal/store/public_execution_test.go @@ -64,7 +64,7 @@ func TestExecutionWorkerAdmissionBindingAndRecovery(t *testing.T) { if inputTextForTest(t, prompt.Input) != "First\n\nSecond" || !prompt.DisableExecutionEnvironment || !prompt.DisableSubagents || prompt.ExecutionControls == nil || *prompt.ExecutionControls != (proto.ExecutionControls{WebSearch: "disabled", TextVerbosity: "medium"}) || prompt.AgentOptions["web_search"] != nil || prompt.AgentOptions["model_verbosity"] != nil { t.Fatal(prompt) } - bound, err := h.s.GetSessionDevice(ctx, h.tenant, h.session.ID) + bound, err := fixtureSessionStore(h.db).GetSessionDevice(ctx, h.tenant, h.session.ID) if err != nil || bound.ID != h.device.ID { t.Fatal(bound, err) } diff --git a/services/core/internal/store/public_handler_fixture_test.go b/services/core/internal/store/public_handler_fixture_test.go index 04a15361d..2dc3288c8 100644 --- a/services/core/internal/store/public_handler_fixture_test.go +++ b/services/core/internal/store/public_handler_fixture_test.go @@ -37,10 +37,11 @@ const testExecutorURL = "wss://core.example/api/v1/agent-daemon/ws" // publicHandler serves s through api.NewHandler. s backs every area the Store // implements, and db is the database and credential key that built s; the -// audit reads, Agents, Files, Vaults, Items, Subagents and Artifacts come from -// db. keys authenticate as Project keys and "admin" as the Core key. Metrics, -// Runtime observation and history, and executor connections are strict -// stand-ins. Execution and Sandboxes stay disabled unless configure sets them. +// audit reads, Agents, Files, Vaults, Environment reads, Items, Subagents and +// Artifacts come from db. keys authenticate as Project keys and "admin" as the +// Core key. Metrics, Runtime observation and history, and executor connections +// are strict stand-ins. Execution and Sandboxes stay disabled unless configure +// sets them. func publicHandler(t testing.TB, s *store.Store, db fixtureDB, keys fixtureKeyResolver, engine string, configure ...func(*api.Dependencies)) (http.Handler, error) { t.Helper() admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential("admin")}) @@ -92,7 +93,7 @@ func publicHandler(t testing.TB, s *store.Store, db fixtureDB, keys fixtureKeyRe Subagents: sessionStore, Artifacts: sessionService, ArtifactsReader: sessionStore, - SessionAdmin: s, Environments: s, Admin: s, AdminAudit: audit, WriteAudit: audit, + SessionAdmin: s, Environments: s, EnvironmentsReader: sessionStore, Admin: s, AdminAudit: audit, WriteAudit: audit, ExecutorConnections: strict, Metrics: strict, RuntimeObservations: strict, RuntimeHistory: strict, } for _, c := range configure { diff --git a/services/core/internal/store/runtime_allocations.go b/services/core/internal/store/runtime_allocations.go index 5a9e45fa9..1979934cc 100644 --- a/services/core/internal/store/runtime_allocations.go +++ b/services/core/internal/store/runtime_allocations.go @@ -56,7 +56,7 @@ func (s *Store) ReserveRuntimeAllocation(ctx context.Context, tenant, environmen if err != nil { return RuntimeAllocation{}, err } - owned, err := s.GetEnvironment(ctx, tenant, environment) + owned, err := sessionpg.LoadEnvironment(ctx, s.queries, tenant, environment) if err != nil { return RuntimeAllocation{}, err } @@ -70,10 +70,11 @@ func (s *Store) ReserveRuntimeAllocation(ctx context.Context, tenant, environmen if err != nil { return RuntimeAllocation{}, err } - device, err := newDeviceParams(lookup.TenantID, "managed-runtime", credentialHash) + 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}) @@ -115,13 +116,13 @@ func (s *Store) ReserveRuntimeAllocation(ctx context.Context, tenant, environmen nodeID = placement.NodeID generation = placement.DeploymentGeneration } - dedicated := sessions.ExecutionDevice{ID: uuid.UUID(device.ID.Bytes).String(), Name: device.Name, EnvironmentID: owned.ID} - if err := sessions.CreateEnvironmentDevice(ctx, sessionpg.BindSession(q, lookup.TenantID, session), dedicated, device.CredentialHash.String); err != nil { + 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: device.ID, ProviderKey: provider, NodeID: nodeID, DeploymentGeneration: generation, + DeviceID: deviceID, ProviderKey: provider, NodeID: nodeID, DeploymentGeneration: generation, }) if err == nil { result = runtimeAllocationFromRow(row, session, lookup.TenantID, pgtype.Timestamptz{}, false) diff --git a/services/core/internal/store/runtime_allocations_test.go b/services/core/internal/store/runtime_allocations_test.go index 29929f859..1e2af94d5 100644 --- a/services/core/internal/store/runtime_allocations_test.go +++ b/services/core/internal/store/runtime_allocations_test.go @@ -23,7 +23,7 @@ func TestRuntimeAllocationAtomicOwnershipAndRecovery(t *testing.T) { if err != nil || owner.Replayed || owner.State != "creating" || owner.CreateSettled { t.Fatalf("reservation: %+v %v", owner, err) } - bound, err := s.GetSessionDevice(t.Context(), tenant, session.ID) + bound, err := sessionAdapter(s).GetSessionDevice(t.Context(), tenant, session.ID) if err != nil || bound.ID != owner.DeviceID || bound.EnvironmentID != environment.ID { t.Fatalf("binding not committed with allocation: %+v %v", bound, err) } @@ -47,7 +47,7 @@ func TestRuntimeAllocationAtomicOwnershipAndRecovery(t *testing.T) { if err != nil || !retry.Replayed || retry.ID != owner.ID || retry.DeviceID != owner.DeviceID { t.Fatalf("restart replaced unknown allocation: %+v %v", retry, err) } - credential, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID) + credential, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID) if err != nil || !ok || credential.CredentialHash != runtimedevice.HashCredential(secret) { t.Fatal("retry rewrote bootstrap credential") } @@ -179,7 +179,7 @@ func TestRuntimeAllocationExpiryAndRevocation(t *testing.T) { if _, err := w.RequestRuntimeCleanup(t.Context(), owner); err != nil { t.Fatal(err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + 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 { diff --git a/services/core/internal/store/runtime_capabilities_pending_test.go b/services/core/internal/store/runtime_capabilities_pending_test.go index 55e7a1d48..9ab5f5fa8 100644 --- a/services/core/internal/store/runtime_capabilities_pending_test.go +++ b/services/core/internal/store/runtime_capabilities_pending_test.go @@ -30,7 +30,7 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { if err != nil { t.Fatal(err) } - env, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + env, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -38,7 +38,7 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { key := uuid.NewString() worker, _ := managedWorkerMode(t, s, db, key, provider, false, true) owner, err := worker.ProvisionEnvironment(t.Context(), tenant, env.ID, key) - if err != nil || initializationState(t, s, owner.TenantID, owner.EnvironmentID) != "pending" { + if err != nil || initializationState(t, db.pool, owner.TenantID, owner.EnvironmentID) != "pending" { t.Fatal(owner, err) } for range 4 { @@ -47,7 +47,7 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { } } owner, err = s.GetRuntimeAllocation(t.Context(), tenant, env.ID) - if err != nil || initializationState(t, s, owner.TenantID, owner.EnvironmentID) != "pending" || owner.State != "running" || provider.writes.Load() != 0 || provider.kills != 0 { + 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) } if _, err := s.GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { diff --git a/services/core/internal/store/runtime_compute_lifecycle_test.go b/services/core/internal/store/runtime_compute_lifecycle_test.go index e40085f4d..2f35bc16f 100644 --- a/services/core/internal/store/runtime_compute_lifecycle_test.go +++ b/services/core/internal/store/runtime_compute_lifecycle_test.go @@ -272,7 +272,7 @@ func newComputeLifecycleFixture(t *testing.T, maxActive, maxRetained int) *compu s, db := newManagedTestStoreDB(t) registry := runtimegateway.NewRegistry() p := &fakeCheckpointProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, computes: map[string]sandbox.ComputeState{}, snapshots: map[string]sandbox.SnapshotIdentity{}, bootstraps: map[string]sandbox.Bootstrap{}, peers: map[string]*websocket.Conn{}, registry: registry} - handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(s), Registry: registry}) + handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(fixtureSessionStore(db)), Registry: registry}) server := httptest.NewServer(http.HandlerFunc(handler.WS)) p.endpoint = "ws" + strings.TrimPrefix(server.URL, "http") t.Cleanup(func() { @@ -307,7 +307,7 @@ func (f *computeLifecycleFixture) sql(query string, args ...any) { func (f *computeLifecycleFixture) create() (string, sessions.Session, sessions.Environment, store.RuntimeAllocation) { t := f.t t.Helper() - tenant, session, environment := managedSession(t, f.store) + tenant, session, environment := managedSession(t, f.store, f.db) owner, err := f.worker.ProvisionEnvironment(t.Context(), tenant, environment.ID, f.key) if err != nil { t.Fatal(err) @@ -503,7 +503,7 @@ func TestRuntimeComputeLifecycleSuspendedDeletionAndExpiryCleanup(t *testing.T) if len(f.provider.computes) != 0 || len(f.provider.snapshots) != 0 || f.provider.snapshotDeletes != 1 { t.Fatal("retained snapshot survived cleanup") } - if _, ok, err := f.store.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := fixtureSessionStore(f.db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("cleanup retained daemon authority", err) } }) @@ -513,7 +513,7 @@ func TestRuntimeComputeLifecycleSuspendedDeletionAndExpiryCleanup(t *testing.T) func TestRuntimeComputeLifecycleCapacityBoundsActiveAndRetained(t *testing.T) { f := newComputeLifecycleFixture(t, 1, 2) tenant, _, env, owner := f.create() - tenant2, _, env2 := managedSession(t, f.store) + tenant2, _, env2 := managedSession(t, f.store, f.db) if _, err := f.worker.ProvisionEnvironment(t.Context(), tenant2, env2.ID, f.key); !errors.Is(err, execution.ErrExecutionUnavailable) { t.Fatalf("active capacity ignored: %v", err) } @@ -545,7 +545,7 @@ func TestRuntimeComputeLifecycleCapacityBoundsActiveAndRetained(t *testing.T) { f.complete(second) f.phase(tenant2, env2.ID, "suspended") // Both retained allocations count even when their source VMs are gone. - tenant3, _, env3 := managedSession(t, f.store) + tenant3, _, env3 := managedSession(t, f.store, f.db) if _, err := f.worker.ProvisionEnvironment(t.Context(), tenant3, env3.ID, f.key); !errors.Is(err, execution.ErrExecutionUnavailable) { t.Fatalf("retained capacity ignored: %v", err) } diff --git a/services/core/internal/store/runtime_configuration_cleanup_test.go b/services/core/internal/store/runtime_configuration_cleanup_test.go index 242d98af7..d188d63bc 100644 --- a/services/core/internal/store/runtime_configuration_cleanup_test.go +++ b/services/core/internal/store/runtime_configuration_cleanup_test.go @@ -90,7 +90,7 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { inspectionError: test.inspectionError, killError: test.killError, } w, _ := managedWorker(t, s, db, key, p) - tenant, session, environment := managedSession(t, s) + tenant, session, environment := managedSession(t, s, db) owner, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key) if (err != nil) != (test.rejectCreate || test.loseCreate) || owner.ID == "" { t.Fatal("unexpected creation outcome", owner, err) @@ -137,7 +137,7 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { if err != nil || got.ID != owner.ID || got.State != wantState || got.CreateSettled != test.wantSettled { t.Fatal("cleanup lost ownership or settlement", got, err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("cleanup retained execution authority", err) } p.mu.Lock() diff --git a/services/core/internal/store/runtime_connection_test.go b/services/core/internal/store/runtime_connection_test.go index 965195511..ab4fa9049 100644 --- a/services/core/internal/store/runtime_connection_test.go +++ b/services/core/internal/store/runtime_connection_test.go @@ -21,10 +21,10 @@ import ( func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { s, db := newManagedTestStoreDB(t) - tenant, session, environment := managedSession(t, s) + tenant, session, environment := managedSession(t, s, db) server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - handler, registry, err := runtime.NewGateway(s, wsURL) + handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), s, wsURL) if err != nil { t.Fatal(err) } @@ -54,7 +54,7 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { if err := w.ReconcileManagedRuntimes(t.Context()); err != nil { t.Fatal(err) } - got, err := s.GetEnvironment(t.Context(), tenant, environment.ID) + got, err := fixtureSessionStore(db).GetEnvironment(t.Context(), tenant, environment.ID) if err != nil { t.Fatal(err) } @@ -90,7 +90,7 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { if err != nil || got.LastTurn != nil || got.EnvironmentInputActivity != nil { t.Fatal("connection fabricated native execution", err) } - if _, err := s.GetEnvironment(t.Context(), uuid.NewString(), environment.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := fixtureSessionStore(db).GetEnvironment(t.Context(), uuid.NewString(), environment.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign Environment access", err) } p.unavailable = true diff --git a/services/core/internal/store/runtime_creation_settlement_test.go b/services/core/internal/store/runtime_creation_settlement_test.go index 6e1460d9a..4a3560c66 100644 --- a/services/core/internal/store/runtime_creation_settlement_test.go +++ b/services/core/internal/store/runtime_creation_settlement_test.go @@ -43,7 +43,7 @@ func TestManagedRuntimeConfirmedAbsentCreateReleasesAtomically(t *testing.T) { key := uuid.NewString() p := &absentCreationProvider{} w, _ := managedWorker(t, s, db, key, p) - tenant, session, environment := managedSession(t, s) + tenant, session, environment := managedSession(t, s, db) ctx, cancel := context.WithCancel(t.Context()) defer cancel() if cancelled { @@ -60,7 +60,7 @@ func TestManagedRuntimeConfirmedAbsentCreateReleasesAtomically(t *testing.T) { if err != nil || stored.State != "released" || !stored.CreateSettled { t.Fatal("release not durable", err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("released credential retained authority", err) } value, err := s.GetSession(t.Context(), tenant, session.ID) @@ -81,7 +81,7 @@ func TestManagedRuntimeForeignAbsenceCannotReleaseCreation(t *testing.T) { key := uuid.NewString() p := &absentCreationProvider{foreign: true} w, _ := managedWorker(t, s, db, key, p) - tenant, _, environment := managedSession(t, s) + tenant, _, environment := managedSession(t, s, db) if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err == nil { t.Fatal("foreign absence accepted") } @@ -95,7 +95,7 @@ func TestManagedRuntimeObservedSettlementAllowsOwnedCleanup(t *testing.T) { key := uuid.NewString() p := &absentCreationProvider{observeSettled: true} w, _ := managedWorker(t, s, db, key, p) - tenant, session, environment := managedSession(t, s) + tenant, session, environment := managedSession(t, s, db) if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err == nil { t.Fatal("uncertain Create succeeded") } diff --git a/services/core/internal/store/runtime_deployment_worker_test.go b/services/core/internal/store/runtime_deployment_worker_test.go index 3570c6cb3..7306de12f 100644 --- a/services/core/internal/store/runtime_deployment_worker_test.go +++ b/services/core/internal/store/runtime_deployment_worker_test.go @@ -15,7 +15,7 @@ func TestManagedDeploymentStartupRejectsSwitchBeforeBackendAccess(t *testing.T) key := uuid.NewString() old := &lifecycleProvider{resources: map[string]sandbox.Info{}} worker, stop := managedWorker(t, s, db, key, old) - tenant, _, environment := managedSession(t, s) + tenant, _, environment := managedSession(t, s, db) owner, err := worker.ProvisionEnvironment(t.Context(), tenant, environment.ID, key) if err != nil { t.Fatal(err) diff --git a/services/core/internal/store/runtime_enrollment.go b/services/core/internal/store/runtime_enrollment.go deleted file mode 100644 index 7ec1dbdd7..000000000 --- a/services/core/internal/store/runtime_enrollment.go +++ /dev/null @@ -1,91 +0,0 @@ -package store - -import ( - "context" - "errors" - - "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/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -// RuntimeEnrollment returns the immutable resource binding, never another secret. -type RuntimeEnrollment struct { - DeviceID string - SessionID string - EnvironmentID string - WorkspaceDirectory string -} - -// EnrollRuntime binds our daemon to one self-hosted Environment. Retries retain -// the same device/key; they cannot replace compute or adopt another native history. -func (s *Store) EnrollRuntime(ctx context.Context, environmentID, credentialHash string) (RuntimeEnrollment, error) { - tenant, err := s.AuthenticateEnvironmentExecutor(ctx, environmentID, credentialHash) - if err != nil { - return RuntimeEnrollment{}, err - } - environment, err := s.GetEnvironment(ctx, tenant, environmentID) - if err != nil { - return RuntimeEnrollment{}, err - } - lookup, err := sessionpg.DeviceLookup(tenant, environmentID) - if err != nil { - return RuntimeEnrollment{}, err - } - var result RuntimeEnrollment - err = s.withPublicSession(ctx, tenant, environment.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - // Recheck authority while holding the Session lock, and retain the key - // lock until binding commits so revocation cannot race enrollment. - authority, err := q.AuthorizeRuntimeEnrollment(ctx, sqlc.AuthorizeRuntimeEnrollmentParams{ - EnvironmentID: lookup.ID, TenantID: lookup.TenantID, TokenSha256: credentialHash, - }) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - bound, err := q.EnrollRuntimeDevice(ctx, sqlc.EnrollRuntimeDeviceParams{ - ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, TenantID: lookup.TenantID, - EnvironmentID: lookup.ID, ExecutorKeyID: authority.KeyID, - }) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrDeviceBindingConflict - } - if err != nil { - return err - } - _, err = q.BindSessionDevice(ctx, sqlc.BindSessionDeviceParams{TenantID: lookup.TenantID, ID: session, ID_2: bound.ID}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrDeviceBindingConflict - } - if err != nil { - return err - } - result = RuntimeEnrollment{DeviceID: uuid.UUID(bound.ID.Bytes).String(), SessionID: environment.SessionID, - EnvironmentID: environment.ID, WorkspaceDirectory: authority.WorkspaceDirectory} - return nil - }) - return result, err -} - -// EnrolledRuntimeBinding identifies user-managed compute, without an allocation -// or any promise of live authorization. The Worker rechecks the socket's key. -type EnrolledRuntimeBinding struct { - DeviceID, TenantID, EnvironmentID, SessionID string -} - -func (s *Store) ListEnrolledRuntimeBindings(ctx context.Context) ([]EnrolledRuntimeBinding, error) { - rows, err := s.queries.ListEnrolledRuntimeBindings(ctx) - if err != nil { - return nil, err - } - result := make([]EnrolledRuntimeBinding, 0, len(rows)) - for _, row := range rows { - result = append(result, EnrolledRuntimeBinding{uuid.UUID(row.DeviceID.Bytes).String(), uuid.UUID(row.TenantID.Bytes).String(), uuid.UUID(row.EnvironmentID.Bytes).String(), uuid.UUID(row.SessionID.Bytes).String()}) - } - return result, nil -} diff --git a/services/core/internal/store/runtime_enrollment_connection_test.go b/services/core/internal/store/runtime_enrollment_connection_test.go index f8d57ecb9..99d805920 100644 --- a/services/core/internal/store/runtime_enrollment_connection_test.go +++ b/services/core/internal/store/runtime_enrollment_connection_test.go @@ -32,7 +32,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) + environment, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), principal.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -41,17 +41,17 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { t.Fatal(err) } digest := sha256.Sum256([]byte(key.Token)) - bound, err := s.EnrollRuntime(t.Context(), environment.ID, hex.EncodeToString(digest[:])) + bound, err := fixtureSessionService(t, db).EnrollRuntime(t.Context(), environment.ID, hex.EncodeToString(digest[:])) if err != nil { t.Fatal(err) } server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - handler, registry, err := runtime.NewGateway(s, wsURL) + handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), s, wsURL) if err != nil { t.Fatal(err) } - connection := runtimeenrollment.ConnectionHandler(s, registry) + connection := runtimeenrollment.ConnectionHandler(fixtureSessionStore(db), registry) assertConnection := func(target, token, status string, code int) { t.Helper() request := httptest.NewRequest("GET", "/api/v1/agent-daemon/connection?environment_id="+target, nil) @@ -121,7 +121,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { t.Helper() deadline := time.Now().Add(5 * time.Second) for time.Now().Before(deadline) { - current, err := s.GetEnvironment(t.Context(), principal.TenantID, environment.ID) + current, err := fixtureSessionStore(db).GetEnvironment(t.Context(), principal.TenantID, environment.ID) if err == nil && current.Status == status { return } diff --git a/services/core/internal/store/runtime_enrollment_test.go b/services/core/internal/store/runtime_enrollment_test.go index bc576a3b4..8384f2f1a 100644 --- a/services/core/internal/store/runtime_enrollment_test.go +++ b/services/core/internal/store/runtime_enrollment_test.go @@ -20,7 +20,7 @@ func runtimeEnrollmentFixture(t *testing.T, s *Store, p identity.Principal) (ses if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), p.TenantID, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), p.TenantID, session.ID) if err != nil { t.Fatal(err) } @@ -42,21 +42,21 @@ func TestRuntimeEnrollmentAuthorityAndRotation(t *testing.T) { p, } { _, target, _ := runtimeEnrollmentFixture(t, s, other) - if _, err := s.EnrollRuntime(ctx, target.ID, executorDigest(key.Token)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionService(t, s).EnrollRuntime(ctx, target.ID, executorDigest(key.Token)); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("foreign target enrollment: %v", err) } } - bound, err := s.EnrollRuntime(ctx, environment.ID, executorDigest(key.Token)) + bound, err := sessionService(t, s).EnrollRuntime(ctx, environment.ID, executorDigest(key.Token)) if err != nil || bound.SessionID != session.ID || bound.EnvironmentID != environment.ID || bound.WorkspaceDirectory != "/workspace" { t.Fatalf("enrollment: %+v %v", bound, err) } - if again, err := s.EnrollRuntime(ctx, environment.ID, executorDigest(key.Token)); err != nil || again != bound { + if again, err := sessionService(t, s).EnrollRuntime(ctx, environment.ID, executorDigest(key.Token)); err != nil || again != bound { t.Fatalf("retry changed binding: %+v %v", again, err) } if devices, err := s.ListExecutionDevices(ctx, p.TenantID); err != nil || len(devices) != 0 { t.Fatalf("enrolled Runtime entered general selection: %v", err) } - auth := runtimegateway.NewAuthenticator(s) + auth := runtimegateway.NewAuthenticator(sessionAdapter(s)) if _, err := auth.AuthenticateBearer(ctx, bound.DeviceID, key.Token); err != nil { t.Fatal(err) } @@ -64,7 +64,7 @@ func TestRuntimeEnrollmentAuthorityAndRotation(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.EnrollRuntime(ctx, environment.ID, executorDigest(otherKey.Token)); !errors.Is(err, sessions.ErrDeviceBindingConflict) { + if _, err := sessionService(t, s).EnrollRuntime(ctx, environment.ID, executorDigest(otherKey.Token)); !errors.Is(err, sessions.ErrDeviceBindingConflict) { t.Fatalf("another key replaced binding: %v", err) } rotated, err := s.RotateExecutorCredential(ctx, p, key.KeyID) @@ -81,21 +81,21 @@ func TestRuntimeEnrollmentAuthorityAndRotation(t *testing.T) { token string denied bool }{{key.Token, true}, {rotated.Token, false}} { - status, err := s.TouchAgentDaemonHeartbeat(ctx, runtimedevice.Heartbeat{RuntimeID: bound.DeviceID, CredentialHash: executorDigest(check.token)}) + status, err := sessionService(t, s).TouchAgentDaemonHeartbeat(ctx, runtimedevice.Heartbeat{RuntimeID: bound.DeviceID, CredentialHash: executorDigest(check.token)}) if err != nil || status.Deleted != check.denied { t.Fatalf("rotation heartbeat: %+v %v", status, err) } } - if again, err := s.EnrollRuntime(ctx, environment.ID, executorDigest(rotated.Token)); err != nil || again != bound { + if again, err := sessionService(t, s).EnrollRuntime(ctx, environment.ID, executorDigest(rotated.Token)); err != nil || again != bound { t.Fatalf("rotation replaced identity: %+v %v", again, err) } if err := s.RevokeExecutorCredential(ctx, p, key.KeyID); err != nil { t.Fatal(err) } - if _, ok, err := s.GetDeviceCredential(ctx, bound.DeviceID); err != nil || ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(ctx, bound.DeviceID); err != nil || ok { t.Fatalf("revoked key authenticates: %v", err) } - if _, err := s.GetSessionDevice(ctx, p.TenantID, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSessionDevice(ctx, p.TenantID, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("revoked binding dispatchable: %v", err) } } @@ -105,13 +105,13 @@ func TestRuntimeEnrollmentConcurrentAndDeletion(t *testing.T) { p := FixtureExecutorPrincipal(t, s, uuid.NewString()) session, environment, key := runtimeEnrollmentFixture(t, s, p) var wg sync.WaitGroup - results := make(chan RuntimeEnrollment, 6) + results := make(chan sessions.RuntimeEnrollment, 6) failures := make(chan error, 6) for range 6 { wg.Add(1) go func() { defer wg.Done() - v, e := s.EnrollRuntime(t.Context(), environment.ID, executorDigest(key.Token)) + v, e := sessionService(t, s).EnrollRuntime(t.Context(), environment.ID, executorDigest(key.Token)) results <- v failures <- e }() @@ -124,7 +124,7 @@ func TestRuntimeEnrollmentConcurrentAndDeletion(t *testing.T) { t.Fatal(err) } } - var bound RuntimeEnrollment + var bound sessions.RuntimeEnrollment for got := range results { if bound.DeviceID == "" { bound = got @@ -140,13 +140,13 @@ func TestRuntimeEnrollmentConcurrentAndDeletion(t *testing.T) { if err := s.DeleteSession(t.Context(), p.TenantID, session.ID); err != nil { t.Fatal(err) } - if _, err := s.EnrollRuntime(t.Context(), environment.ID, executorDigest(key.Token)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionService(t, s).EnrollRuntime(t.Context(), environment.ID, executorDigest(key.Token)); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("deleted enrollment: %v", err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), bound.DeviceID); err != nil || ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), bound.DeviceID); err != nil || ok { t.Fatalf("deleted Session authenticates: %v", err) } - status, err := s.TouchAgentDaemonHeartbeat(t.Context(), runtimedevice.Heartbeat{RuntimeID: bound.DeviceID, CredentialHash: executorDigest(key.Token)}) + status, err := sessionService(t, s).TouchAgentDaemonHeartbeat(t.Context(), runtimedevice.Heartbeat{RuntimeID: bound.DeviceID, CredentialHash: executorDigest(key.Token)}) if err != nil || !status.Deleted { t.Fatalf("deleted heartbeat: %+v %v", status, err) } diff --git a/services/core/internal/store/runtime_environment_terminal_test.go b/services/core/internal/store/runtime_environment_terminal_test.go index 4249b354c..e7f545584 100644 --- a/services/core/internal/store/runtime_environment_terminal_test.go +++ b/services/core/internal/store/runtime_environment_terminal_test.go @@ -49,7 +49,7 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin if failure := ended.EnvironmentFailure; expired != (failure == nil) || !expired && (failure.Reason != sessions.ProvisioningFailureReason || failure.FailedAt.IsZero()) { t.Fatal("terminal failure projection", failure) } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("terminal credential remained usable", err) } failed, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) @@ -99,7 +99,7 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin if next, err := s.SessionEventCursor(t.Context(), tenant, session.ID); err != nil || next != cursor { t.Fatal("cleanup repeated terminal events", next, err) } - if err := writer.ReplaceEnvironmentConnection(t.Context(), tenant, session.Environment.ID, uuid.NewString()); !errors.Is(err, sessions.ErrInvalidInput) { + 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) { @@ -145,7 +145,7 @@ func TestManagedEnvironmentFailureRollsBackWithSessionEvent(t *testing.T) { if err != nil || allocation.State != "creating" { t.Fatal("partial allocation transition", err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || !ok { + if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || !ok { t.Fatal("partial credential revocation", err) } } diff --git a/services/core/internal/store/runtime_history_acceptance_test.go b/services/core/internal/store/runtime_history_acceptance_test.go index 1334b9e45..fb496e743 100644 --- a/services/core/internal/store/runtime_history_acceptance_test.go +++ b/services/core/internal/store/runtime_history_acceptance_test.go @@ -28,14 +28,14 @@ func historyBackend(t *testing.T, s *store.Store) *postgresreader.Reader { } return r } -func historyOwner(t *testing.T, s *store.Store) runtimehistory.Scope { +func historyOwner(t *testing.T, s *store.Store, pool *pgxpool.Pool) runtimehistory.Scope { t.Helper() tenant := uuid.NewString() session, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"model":"fixture-model"},"environment":{"type":"openai_hosted","workspace_directory":"/workspace","capability_directories":[]}}`)}) if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionReads(pool).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -57,8 +57,8 @@ func historyQuery(scope runtimehistory.Scope, start, end time.Time, points int) func TestPostgresRuntimeHistoryAcceptance(t *testing.T) { s, pool := store.NewTestStore(t) - scope := historyOwner(t, s) - foreign := historyOwner(t, s) + scope := historyOwner(t, s, pool) + foreign := historyOwner(t, s, pool) reader := historyBackend(t, s) end := time.Now().UTC().Truncate(time.Second) start := end.Add(-3 * time.Minute) @@ -138,7 +138,7 @@ func TestPostgresRuntimeHistoryAcceptance(t *testing.T) { func TestPostgresRuntimeHistoryDenseReadAndBoundedRetention(t *testing.T) { s, pool := store.NewTestStore(t) - scope := historyOwner(t, s) + scope := historyOwner(t, s, pool) reader := historyBackend(t, s) end := time.Now().UTC().Truncate(time.Second) start := end.Add(-24 * time.Hour) @@ -206,8 +206,8 @@ func TestPostgresRuntimeHistoryDenseReadAndBoundedRetention(t *testing.T) { } func TestPostgresRuntimeHistoryKeepsProviderReportedUtilization(t *testing.T) { - s, _ := store.NewTestStore(t) - scope := historyOwner(t, s) + s, pool := store.NewTestStore(t) + scope := historyOwner(t, s, pool) reader := historyBackend(t, s) allocation := uuid.NewString() start := time.Now().UTC().Truncate(time.Minute).Add(-10 * time.Minute) diff --git a/services/core/internal/store/runtime_initialization_test.go b/services/core/internal/store/runtime_initialization_test.go index 5aefa29ca..e56fa37a1 100644 --- a/services/core/internal/store/runtime_initialization_test.go +++ b/services/core/internal/store/runtime_initialization_test.go @@ -6,13 +6,14 @@ import ( "context" "encoding/json" "errors" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentplugin" "reflect" "strings" "sync" "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/internal/agentplugin" + 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" @@ -62,7 +63,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { if err != nil { t.Fatal(err) } - env, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + env, err := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -73,7 +74,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { } p := &initializingProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, initializationPeer: initializationPeer{deferred: true}} p.apply = func(_ proto.RuntimePreparePayload, _ []byte) proto.RuntimePrepareResultPayload { - if _, err := s.GetSessionDevice(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := fixtureSessionStore(db).GetSessionDevice(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Error("premature file access", err) } if _, err := s.GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { @@ -87,7 +88,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { key := uuid.NewString() w, stop := managedWorkerMode(t, s, db, key, p, false, true) if mode == "restart" { - awaitInitialization(t, s, tenant, env.ID, "failed") + awaitInitialization(t, db.pool, tenant, env.ID, "failed") if p.writes.Load() != 0 { t.Fatal("recovered unknown operation replayed") } @@ -97,11 +98,11 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { if err != nil { t.Fatal(err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || !ok { + if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || !ok { t.Fatal("preparation blocked authentication", err) } time.Sleep(350 * time.Millisecond) - if initializationState(t, s, tenant, env.ID) != "pending" || p.writes.Load() != 0 { + if initializationState(t, db.pool, tenant, env.ID) != "pending" || p.writes.Load() != 0 { t.Fatal("missing socket consumed initialization") } p.deferred = false @@ -112,7 +113,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { if mode == "uncertain" { want = "failed" } - awaitInitialization(t, s, tenant, env.ID, want) + awaitInitialization(t, db.pool, tenant, env.ID, want) if mode == "complete" { if int(p.writes.Load()) != expectedSteps { t.Fatal("missing operations", p.writes.Load()) @@ -157,7 +158,7 @@ func TestManagedRuntimePreparationAllOperationsUsePeer(t *testing.T) { } tenant := uuid.NewString() fileBody := bytes.Repeat([]byte("bounded bytes"), 12000) - session, environment := hostedFailureSession(t, s, tenant, sessions.CreateSession{ + session, environment := hostedFailureSession(t, s, db, tenant, sessions.CreateSession{ InitialFiles: []environmentconfig.InitialFile{{Type: "inline", Path: "/workspace/first", Data: fileBody}}, Initialization: environmentconfig.Setup{Skills: []environmentconfig.Skill{hostedFailureSkill(t)}, Plugins: []environmentconfig.Plugin{{Metadata: agentplugin.Metadata{Type: "inline", Name: "plugin", Description: "A plugin."}, Archive: archive.Bytes()}}, Packages: v1.EnvironmentPackages{NPM: []string{"is-number@7.0.0"}, Python: []string{"packaging==24.2"}}, Commands: []environmentconfig.SetupCommand{{Command: "read installed bundles and create directory"}}, CapabilityDirectories: []string{"/workspace/generated"}}, }) @@ -185,7 +186,7 @@ func TestManagedRuntimePreparationAllOperationsUsePeer(t *testing.T) { if _, err := worker.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err != nil { t.Fatal(err) } - awaitInitialization(t, s, tenant, environment.ID, "complete") + awaitInitialization(t, db.pool, tenant, environment.ID, "complete") actionsMu.Lock() defer actionsMu.Unlock() expected := []string{"file", "configure", "skill", "plugin", "npm", "python", "setup", "finalize"} @@ -193,7 +194,7 @@ func TestManagedRuntimePreparationAllOperationsUsePeer(t *testing.T) { t.Fatal("typed ordering or provider isolation", actions, provider.commandCalls.Load()) } allocation, err := s.GetRuntimeAllocation(t.Context(), tenant, environment.ID) - if err != nil || initializationState(t, s, allocation.TenantID, allocation.EnvironmentID) != "complete" { + 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 6f093aa19..afef6de5c 100644 --- a/services/core/internal/store/runtime_input_admission_test.go +++ b/services/core/internal/store/runtime_input_admission_test.go @@ -14,7 +14,7 @@ import ( func TestManagedRuntimeMaintenancePreservesCancelAndRetry(t *testing.T) { s, db := newManagedTestStoreDB(t) - tenant, session, _ := managedSession(t, s) + tenant, session, _ := managedSession(t, s, db) inputs := []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"accepted work"}`)}} accepted, err := s.SubmitInputs(t.Context(), tenant, session.ID, "work", inputs) if err != nil { diff --git a/services/core/internal/store/runtime_lifecycle_test.go b/services/core/internal/store/runtime_lifecycle_test.go index 8a6af2daf..c7cd5e3eb 100644 --- a/services/core/internal/store/runtime_lifecycle_test.go +++ b/services/core/internal/store/runtime_lifecycle_test.go @@ -83,7 +83,7 @@ func managedWorkerMode(t *testing.T, s *store.Store, db fixtureDB, key string, p if peer, ok := p.(interface { setRuntimeGateway(*testing.T, string, *runtimegateway.Registry) }); ok { - handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(s), Registry: registry}) + handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(fixtureSessionStore(db)), Registry: registry}) server := httptest.NewServer(http.HandlerFunc(handler.WS)) t.Cleanup(server.Close) peer.setRuntimeGateway(t, "ws"+strings.TrimPrefix(server.URL, "http"), registry) @@ -113,14 +113,14 @@ func managedWorkerMode(t *testing.T, s *store.Store, db fixtureDB, key string, p return w, stop } -func managedSession(t *testing.T, s *store.Store) (string, sessions.Session, sessions.Environment) { +func managedSession(t *testing.T, s *store.Store, db fixtureDB) (string, sessions.Session, sessions.Environment) { t.Helper() tenant := uuid.NewString() v, e := s.CreateSession(t.Context(), tenant, store.WithFixtureModelProvider(sessions.CreateSession{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"model":"test"},"environment":{"type":"openai_hosted","network":{"access":"enabled"}}}`)})) if e != nil { t.Fatal(e) } - env, e := s.GetSessionEnvironment(t.Context(), tenant, v.ID) + env, e := fixtureSessionStore(db).GetSessionEnvironment(t.Context(), tenant, v.ID) if e != nil { t.Fatal(e) } @@ -146,7 +146,7 @@ func reconcileManagedState(t *testing.T, w *execution.Worker, s *store.Store, te func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { s, db := newManagedTestStoreDB(t) - tenant, session, env := managedSession(t, s) + tenant, session, env := managedSession(t, s, db) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}, loseCreate: true} w, stop := managedWorker(t, s, db, key, p) @@ -154,7 +154,7 @@ func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { if err == nil || owner.ID == "" { t.Fatal("fault did not retain allocation") } - credential, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID) + credential, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID) if err != nil || !ok || credential.CredentialHash != p.credentialHash { t.Fatal("provider received unbound credential") } @@ -178,14 +178,14 @@ func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { if err != nil || clean.State != "released" || p.kills != 1 { t.Fatalf("deleted cleanup: %+v %v", clean, err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("cleanup did not revoke authority") } } func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { s, db := newManagedTestStoreDB(t) - tenant, session, env := managedSession(t, s) + tenant, session, env := managedSession(t, s, db) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}, loseCreate: true, absent: true} w, _ := managedWorker(t, s, db, key, p) @@ -201,7 +201,7 @@ func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { if err != nil || got.State != "cleanup_pending" || got.CreateSettled || p.creates != 1 { t.Fatalf("unknown creation forgotten: %+v %v", got, err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("unknown allocation retains execution authority") } // A late completion is still owned and reclaimed on the next scan. @@ -215,7 +215,7 @@ func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { func TestManagedRuntimeExpiryRevokesWhenProviderUnavailable(t *testing.T) { s, db := newManagedTestStoreDB(t) - tenant, _, env := managedSession(t, s) + tenant, _, env := managedSession(t, s, db) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} w, _ := managedWorker(t, s, db, key, p) @@ -232,7 +232,7 @@ func TestManagedRuntimeExpiryRevokesWhenProviderUnavailable(t *testing.T) { if err != nil || got.State != "cleanup_pending" { t.Fatalf("expiry lost on provider failure: %+v %v", got, err) } - if _, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { + if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("expired credential still authenticates") } if p.kills != 0 { @@ -242,7 +242,7 @@ func TestManagedRuntimeExpiryRevokesWhenProviderUnavailable(t *testing.T) { func TestManagedRuntimeStoppedComputeDoesNotRequestCleanup(t *testing.T) { s, db := newManagedTestStoreDB(t) - tenant, _, env := managedSession(t, s) + tenant, _, env := managedSession(t, s, db) key := uuid.NewString() p := &lifecycleProvider{resources: map[string]sandbox.Info{}} w, _ := managedWorker(t, s, db, key, p) 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 da9e9c8d1..23c4776c6 100644 --- a/services/core/internal/store/runtime_node_lifecycle_fixture_test.go +++ b/services/core/internal/store/runtime_node_lifecycle_fixture_test.go @@ -114,7 +114,7 @@ func newNodeIsolationFixture(t *testing.T, mode string) *nodeIsolationFixture { p.writes.Add(1) return completedInitialization(request, data) }} - handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(s), Registry: registry}) + handler := runtimegateway.NewHandler(runtimegateway.HandlerConfig{Authenticator: runtimegateway.NewAuthenticator(fixtureSessionStore(db)), Registry: registry}) server := httptest.NewServer(http.HandlerFunc(handler.WS)) cp.endpoint = "ws" + strings.TrimPrefix(server.URL, "http") t.Cleanup(func() { @@ -210,7 +210,7 @@ func (f *nodeIsolationFixture) session(node string, initialize bool) (string, se if err != nil { f.t.Fatal(err) } - env, err := f.store.GetSessionEnvironment(f.t.Context(), tenant, session.ID) + env, err := sessionReads(f.pool).GetSessionEnvironment(f.t.Context(), tenant, session.ID) if err != nil { f.t.Fatal(err) } diff --git a/services/core/internal/store/runtime_node_lifecycle_test.go b/services/core/internal/store/runtime_node_lifecycle_test.go index ed9295804..8046c574a 100644 --- a/services/core/internal/store/runtime_node_lifecycle_test.go +++ b/services/core/internal/store/runtime_node_lifecycle_test.go @@ -113,8 +113,8 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { f.provider.mu.Lock() restores := f.provider.restores f.provider.mu.Unlock() - state := fmt.Sprintf("wake=%s/%s err=%v; deleted=%s/%s err=%v; initialized=%s/%s/%s err=%v; restores=%d writes=%d blocked_returns=%d", wake.State, wake.ComputePhase, e1, deleted.State, deleted.ComputePhase, e2, initialized.State, initializationState(t, f.store, initialized.TenantID, initialized.EnvironmentID), initialized.ComputePhase, e3, restores, f.provider.writes.Load(), f.provider.returned.Load()) - return e1 == nil && e2 == nil && e3 == nil && wake.ComputePhase == "running" && deleted.State == "released" && initializationState(t, f.store, initialized.TenantID, initialized.EnvironmentID) == "complete" && initialized.ComputePhase == "running", state + state := fmt.Sprintf("wake=%s/%s err=%v; deleted=%s/%s err=%v; initialized=%s/%s/%s err=%v; restores=%d writes=%d blocked_returns=%d", wake.State, wake.ComputePhase, e1, deleted.State, deleted.ComputePhase, e2, initialized.State, initializationState(t, f.pool, initialized.TenantID, initialized.EnvironmentID), initialized.ComputePhase, e3, restores, f.provider.writes.Load(), f.provider.returned.Load()) + return e1 == nil && e2 == nil && e3 == nil && wake.ComputePhase == "running" && deleted.State == "released" && initializationState(t, f.pool, initialized.TenantID, initialized.EnvironmentID) == "complete" && initialized.ComputePhase == "running", state }) if f.provider.returned.Load() != 0 || f.provider.writes.Load() != 1 { t.Fatalf("A returned early or initialization replayed: returned=%d writes=%d", f.provider.returned.Load(), f.provider.writes.Load()) diff --git a/services/core/internal/store/runtime_observation_scan_test.go b/services/core/internal/store/runtime_observation_scan_test.go index 3b00e6532..4aca54d61 100644 --- a/services/core/internal/store/runtime_observation_scan_test.go +++ b/services/core/internal/store/runtime_observation_scan_test.go @@ -8,10 +8,10 @@ import ( ) func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *testing.T) { - s, _ := store.NewManagedTestStore(t) + s, db := newManagedTestStoreDB(t) var expected []string for range 5 { - _, session, _ := managedSession(t, s) + _, session, _ := managedSession(t, s, db) expected = append(expected, session.ID) } slices.Sort(expected) diff --git a/services/core/internal/store/runtime_pending_test.go b/services/core/internal/store/runtime_pending_test.go index 0aa82eca0..f21831945 100644 --- a/services/core/internal/store/runtime_pending_test.go +++ b/services/core/internal/store/runtime_pending_test.go @@ -15,12 +15,12 @@ import ( func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) { s, db := newManagedTestStoreDB(t) - tenant, idle, idleEnvironment := managedSession(t, s) + tenant, idle, idleEnvironment := managedSession(t, s, db) initial, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"model":"test"},"environment":{"type":"openai_hosted"}}`), InitialInputs: []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"hello"}`)}}}) if err != nil { t.Fatal(err) } - _, deleted, deletedEnvironment := managedSession(t, s) + _, deleted, deletedEnvironment := managedSession(t, s, db) if err := s.DeleteSession(t.Context(), deleted.TenantID, deleted.ID); err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/runtime_scan_test.go b/services/core/internal/store/runtime_scan_test.go index 7cd13fd66..4845c2b03 100644 --- a/services/core/internal/store/runtime_scan_test.go +++ b/services/core/internal/store/runtime_scan_test.go @@ -32,7 +32,7 @@ func TestManagedRuntimeScanWrapServicesNextPage(t *testing.T) { w, _ := managedWorker(t, s, db, key, p) var ids []string for range count { - tenant, _, env := managedSession(t, s) + tenant, _, env := managedSession(t, s, db) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) @@ -69,7 +69,7 @@ func TestManagedRuntimeScanEmptyAfterCleanupAndCanceledCall(t *testing.T) { p := &scanProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}} key := uuid.NewString() w, _ := managedWorker(t, s, db, key, p) - tenant, session, env := managedSession(t, s) + tenant, session, env := managedSession(t, s, db) owner, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) @@ -108,7 +108,7 @@ func TestManagedRuntimeScanEmptyAfterCleanupAndCanceledCall(t *testing.T) { } } // A new allocation remains discoverable after the store becomes empty. - tenant, _, env = managedSession(t, s) + tenant, _, env = managedSession(t, s, db) next, err := w.ProvisionEnvironment(t.Context(), tenant, env.ID, key) if err != nil { t.Fatal(err) diff --git a/services/core/internal/store/runtime_suspension.go b/services/core/internal/store/runtime_suspension.go index 9b0e2a5f6..da67bd122 100644 --- a/services/core/internal/store/runtime_suspension.go +++ b/services/core/internal/store/runtime_suspension.go @@ -121,7 +121,7 @@ func (s *Store) TouchRuntimeActivity(ctx context.Context, tenant, environment st if err != nil { return err } - owned, err := s.GetEnvironment(ctx, tenant, environment) + owned, err := sessionpg.LoadEnvironment(ctx, s.queries, tenant, environment) if err != nil { return err } diff --git a/services/core/internal/store/runtime_worker_recovery_test.go b/services/core/internal/store/runtime_worker_recovery_test.go index 0edabc0b6..db988287d 100644 --- a/services/core/internal/store/runtime_worker_recovery_test.go +++ b/services/core/internal/store/runtime_worker_recovery_test.go @@ -74,7 +74,7 @@ func TestWorkerWaitsForComputeAndSurvivesPromotionConflict(t *testing.T) { } }) frames := workerFrames(t, h) - environment, err := h.s.GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) + environment, err := fixtureSessionStore(h.db).GetSessionEnvironment(t.Context(), h.tenant, h.session.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/sandbox_deployment_switch_test.go b/services/core/internal/store/sandbox_deployment_switch_test.go index 3fa85be08..a101ea7f5 100644 --- a/services/core/internal/store/sandbox_deployment_switch_test.go +++ b/services/core/internal/store/sandbox_deployment_switch_test.go @@ -118,7 +118,7 @@ func TestSandboxDirectDeploymentOwnershipAndCleanSwitch(t *testing.T) { if err != nil || owner.NodeID != "" { t.Fatal(owner, err) } - credential, ok, err := s.GetDeviceCredential(t.Context(), owner.DeviceID) + credential, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID) if err != nil || !ok || credential.RuntimeAllocationID != owner.ID || credential.RuntimeNodeID != "" { t.Fatal("direct bootstrap lost managed identity", err) } 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 a1c0dd1cb..c9c8a5021 100644 --- a/services/core/internal/store/sandbox_deployment_switch_worker_test.go +++ b/services/core/internal/store/sandbox_deployment_switch_worker_test.go @@ -145,7 +145,7 @@ func TestSandboxWorkerSwitchesAndRecoversFailedActivation(t *testing.T) { if _, err := w.InitializeSandboxDeployment(t.Context(), cloud); err != nil { t.Fatal("setup after unconfigured publication", err) } - tenant, session, environment := managedSession(t, s) + tenant, session, environment := managedSession(t, s, db) allocation, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, id) if err != nil || allocation.NodeID != "" || allocation.State != "running" { t.Fatal("direct provider not available after resume", allocation, err) diff --git a/services/core/internal/store/sandbox_deployment_worker_test.go b/services/core/internal/store/sandbox_deployment_worker_test.go index b797c6274..039441dd4 100644 --- a/services/core/internal/store/sandbox_deployment_worker_test.go +++ b/services/core/internal/store/sandbox_deployment_worker_test.go @@ -73,7 +73,7 @@ func TestSandboxDeploymentWorkerActivatesWithoutRestart(t *testing.T) { } } connect() - tenant, _, environment := managedSession(t, s) + tenant, _, environment := managedSession(t, s, db) allocation, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, id) if err != nil || allocation.NodeID != nodeID || p.creates != 1 { t.Fatal("activation failed", allocation, err) diff --git a/services/core/internal/store/sandbox_reset_test.go b/services/core/internal/store/sandbox_reset_test.go index 93812c870..9c4648e4b 100644 --- a/services/core/internal/store/sandbox_reset_test.go +++ b/services/core/internal/store/sandbox_reset_test.go @@ -307,7 +307,7 @@ func TestSandboxResetSnapshotCountsOfflineOwnershipOnce(t *testing.T) { if err := deploymentService(t, s).RemoveNode(t.Context(), process.LocalNodeID); !errors.Is(err, deployment.ErrNodeInUse) { t.Fatal("removed node with reset resources", err) } - if row, err := s.GetEnvironment(t.Context(), pending.TenantID, pending.Environment.ID); err == nil && row.Status == "expired" { + if row, err := sessionAdapter(s).GetEnvironment(t.Context(), pending.TenantID, pending.Environment.ID); err == nil && row.Status == "expired" { t.Fatal("projection mutated pending resource") } } diff --git a/services/core/internal/store/self_hosted_initial_public_test.go b/services/core/internal/store/self_hosted_initial_public_test.go index 53fc88bb7..90bcb379b 100644 --- a/services/core/internal/store/self_hosted_initial_public_test.go +++ b/services/core/internal/store/self_hosted_initial_public_test.go @@ -130,7 +130,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { if !reflect.DeepEqual(texts, item.Texts) { t.Fatal("public initial text order changed") } - environment, err := s.GetSessionEnvironment(t.Context(), tenant, item.ID) + environment, err := sessionReads(pool).GetSessionEnvironment(t.Context(), tenant, item.ID) if err != nil || environment.ID != item.EnvironmentID || environment.Status != "pending" { t.Fatal("public initial Environment identity changed", err) } diff --git a/services/core/internal/store/session_artifacts_test.go b/services/core/internal/store/session_artifacts_test.go index f37d5e744..2504d2916 100644 --- a/services/core/internal/store/session_artifacts_test.go +++ b/services/core/internal/store/session_artifacts_test.go @@ -55,7 +55,7 @@ func artifactTurn(t *testing.T, s *Store, kind string) (tenant, session, environ if err != nil { t.Fatal(err) } - env, err := s.GetSessionEnvironment(t.Context(), tenant, created.ID) + env, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, created.ID) if err != nil { t.Fatal(err) } @@ -409,7 +409,7 @@ func TestSessionArtifactsRepublishOnlyNewChangedOrDeletedPaths(t *testing.T) { if err != nil { t.Fatal(err) } - otherEnvironment, err := s.GetSessionEnvironment(t.Context(), tenant, other.ID) + otherEnvironment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, other.ID) if err != nil { t.Fatal(err) } @@ -439,7 +439,7 @@ func TestSessionArtifactsNewestVersionFollowsTurnOrder(t *testing.T) { t.Fatal(err) } session := created.ID - env, err := s.GetSessionEnvironment(t.Context(), tenant, session) + env, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_deletion_lifecycle_public_test.go b/services/core/internal/store/session_deletion_lifecycle_public_test.go index b5368fbac..5eaf767e5 100644 --- a/services/core/internal/store/session_deletion_lifecycle_public_test.go +++ b/services/core/internal/store/session_deletion_lifecycle_public_test.go @@ -95,10 +95,10 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { connect := func(session sessions.Session) { t.Helper() generation := uuid.NewString() - if err := writer.ReplaceEnvironmentConnection(ctx, tenant, session.Environment.ID, generation); err != nil { + if err := leased.Sessions.ReplaceEnvironmentConnection(ctx, tenant, session.Environment.ID, generation); err != nil { t.Fatal(err) } - if err := writer.ObserveEnvironmentConnection(ctx, tenant, session.Environment.ID, generation, 1, true); err != nil { + if err := leased.Sessions.ObserveEnvironmentConnection(ctx, tenant, session.Environment.ID, generation, 1, true); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/session_devices_fixture_test.go b/services/core/internal/store/session_devices_fixture_test.go new file mode 100644 index 000000000..582f1197a --- /dev/null +++ b/services/core/internal/store/session_devices_fixture_test.go @@ -0,0 +1,33 @@ +package store_test + +import ( + "context" + "errors" + "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// bindSessionDevice binds the tenant's device to the Session through the +// Session execution operations on an execution lease of its own, and returns +// once the server released that lease, so a Worker can take it next. +func bindSessionDevice(t *testing.T, db fixtureDB, tenant, session, device string) error { + t.Helper() + lease, err := pgunit.AcquireLease(t.Context(), db.pool) + if err != nil { + return err + } + released := pgtest.ObserveExecutionLeaseRelease(t, db.pool) + operations, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) + if err == nil { + err = operations.BindSessionDevice(t.Context(), tenant, session, device) + } + if err := errors.Join(err, lease.Close(context.Background())); err != nil { + return err + } + released() + return nil +} diff --git a/services/core/internal/store/session_diagnostics_test.go b/services/core/internal/store/session_diagnostics_test.go index 9a028592d..582779d87 100644 --- a/services/core/internal/store/session_diagnostics_test.go +++ b/services/core/internal/store/session_diagnostics_test.go @@ -221,7 +221,7 @@ func TestDiagnosticProvisioningDetailAtomicAndPrivate(t *testing.T) { failure := sessions.ProvisioningFailure{Step: sessions.ProvisioningSetupCommand, Index: 2, ExitCode: 7} runtimeSuspensionSQL(t, pool, "UPDATE environments SET initialization='running' WHERE id=$1", owner.EnvironmentID) preparation := sessions.EnvironmentInitialization{EnvironmentID: owner.EnvironmentID, SessionID: owner.SessionID, TenantID: owner.TenantID, DeviceID: owner.DeviceID} - if err = writer.FailEnvironmentInitialization(t.Context(), preparation, failure); err == nil { + if err = sessionExecution(t, writer.lease).FailEnvironmentInitialization(t.Context(), preparation, failure); err == nil { t.Fatal("failure committed without events") } var detail []byte @@ -229,7 +229,7 @@ func TestDiagnosticProvisioningDetailAtomicAndPrivate(t *testing.T) { t.Fatal("partial failure detail", string(detail), err) } runtimeSuspensionSQL(t, pool, "ALTER TABLE session_events DROP CONSTRAINT "+constraint) - if err = writer.FailEnvironmentInitialization(t.Context(), preparation, failure); err != nil { + if err = sessionExecution(t, writer.lease).FailEnvironmentInitialization(t.Context(), preparation, failure); err != nil { t.Fatal(err) } snap, err := s.GetSessionDiagnosticsSnapshot(t.Context(), tenant, session.ID) diff --git a/services/core/internal/store/session_execution_configuration_test.go b/services/core/internal/store/session_execution_configuration_test.go index 0f80f4d0b..e44de2a52 100644 --- a/services/core/internal/store/session_execution_configuration_test.go +++ b/services/core/internal/store/session_execution_configuration_test.go @@ -256,7 +256,7 @@ func TestSessionExecutionConfigurationSurvivesSuspendResume(t *testing.T) { if err != nil { t.Fatal(err) } - environment, err := s.GetSessionEnvironment(t.Context(), tenant, session.ID) + environment, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_initial_input.go b/services/core/internal/store/session_initial_input.go index 18f34a14b..e08ef4486 100644 --- a/services/core/internal/store/session_initial_input.go +++ b/services/core/internal/store/session_initial_input.go @@ -98,9 +98,15 @@ func (s *Store) createSessionResources(ctx context.Context, tenant string, param return err } } - if err := createSessionEnvironment(ctx, q, row); err != nil { + creates, err := sessions.CreatesEnvironment(row.Configuration) + if err != nil { return err } + if creates { + if err := sessionpg.CreateSessionEnvironment(ctx, q, row.ID); err != nil { + return err + } + } if placement.Environment.Type == "openai_hosted" { if err := reserveRuntimePlacement(ctx, q, row.ID, s.publicURL); err != nil { return err diff --git a/services/core/internal/store/session_reads_fixture_test.go b/services/core/internal/store/session_reads_fixture_test.go new file mode 100644 index 000000000..5ddb4be8c --- /dev/null +++ b/services/core/internal/store/session_reads_fixture_test.go @@ -0,0 +1,43 @@ +package store + +import ( + "context" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// FixtureEnvironmentDevice provisions the dedicated Runtime device of the +// tenant's hosted Environment on pool through the procedure the managed +// Runtime allocation runs, without the allocation. +func FixtureEnvironmentDevice(ctx context.Context, pool *pgxpool.Pool, tenant, environment, name, credentialHash string) (sessions.ExecutionDevice, error) { + units := pgunit.NewPool(pool) + current, err := sessionpg.New(units, nil).GetEnvironment(ctx, tenant, environment) + if err != nil { + return sessions.ExecutionDevice{}, err + } + registration, err := sessions.NewDeviceRegistration(name, credentialHash) + if err != nil { + return sessions.ExecutionDevice{}, err + } + lookup, err := sessionpg.DeviceLookup(tenant, current.SessionID) + if err != nil { + return sessions.ExecutionDevice{}, err + } + device := sessions.ExecutionDevice{ID: uuid.NewString(), Name: registration.Name, EnvironmentID: current.ID} + err = sessionpg.WithSession(ctx, units, lookup.TenantID, lookup.ID, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + if err := locked.Public(); err != nil { + return err + } + return sessions.CreateEnvironmentDevice(ctx, sessionpg.BindSession(q, lookup.TenantID, lookup.ID), device, registration.CredentialHash) + }) + if err != nil { + return sessions.ExecutionDevice{}, err + } + return device, nil +} diff --git a/services/core/internal/store/session_write_audit_test.go b/services/core/internal/store/session_write_audit_test.go index 505214b98..c72082ed9 100644 --- a/services/core/internal/store/session_write_audit_test.go +++ b/services/core/internal/store/session_write_audit_test.go @@ -53,7 +53,7 @@ func TestSessionWriteAuditCreationReplayNoopAndDeletion(t *testing.T) { if err != nil { t.Fatal(err) } - env, err := s.GetSessionEnvironment(t.Context(), tenant, created.ID) + env, err := sessionAdapter(s).GetSessionEnvironment(t.Context(), tenant, created.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/sessions_fixture_test.go b/services/core/internal/store/sessions_fixture_test.go index 7b25eb1b1..bd878fbc4 100644 --- a/services/core/internal/store/sessions_fixture_test.go +++ b/services/core/internal/store/sessions_fixture_test.go @@ -1,6 +1,8 @@ package store_test import ( + "testing" + "github.com/jackc/pgx/v5/pgxpool" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" @@ -8,13 +10,32 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -// sessionReads is the Session adapter on pool. It serves the Item, Subagent -// and Artifact reads. -func sessionReads(pool *pgxpool.Pool) *sessionpg.Store { return sessionpg.New(pgunit.NewPool(pool)) } +// sessionReads is the Session adapter on pool without a credential key. It +// serves the Session reads that open no frozen data. +func sessionReads(pool *pgxpool.Pool) *sessionpg.Store { + return sessionpg.New(pgunit.NewPool(pool), nil) +} // fixtureSessions builds the Session adapter and service on db, as cmd/server does. func fixtureSessions(db fixtureDB) (*sessionpg.Store, *sessions.Service, error) { - sessionStore := sessionReads(db.pool) + sessionStore := fixtureSessionStore(db) sessionService, err := sessions.NewService(sessionStore) return sessionStore, sessionService, err } + +// fixtureSessionStore is the Session adapter on db with its credential key, as +// cmd/server builds it. +func fixtureSessionStore(db fixtureDB) *sessionpg.Store { + return sessionpg.New(pgunit.NewPool(db.pool), db.cipher) +} + +// fixtureSessionService is the Session service on fixtureSessionStore(db). It +// runs the device, heartbeat and enrollment use cases. +func fixtureSessionService(t testing.TB, db fixtureDB) *sessions.Service { + t.Helper() + service, err := sessions.NewService(fixtureSessionStore(db)) + if err != nil { + t.Fatal(err) + } + return service +} diff --git a/services/core/internal/store/sessions_test.go b/services/core/internal/store/sessions_test.go index cb3976fe9..8227d2f89 100644 --- a/services/core/internal/store/sessions_test.go +++ b/services/core/internal/store/sessions_test.go @@ -22,9 +22,20 @@ func testStore(t *testing.T) (*Store, *pgxpool.Pool) { return New(pool), pool } -// sessionAdapter is the Session adapter on s's database. It serves the Item, -// Subagent and Artifact reads and deletes Artifacts. -func sessionAdapter(s *Store) *sessionpg.Store { return sessionpg.New(s.pooled) } +// sessionAdapter is the Session adapter on s's database with s's credential +// key. It serves the Session reads and pooled Session storage. +func sessionAdapter(s *Store) *sessionpg.Store { return sessionpg.New(s.pooled, s.credentialCipher) } + +// sessionService is the Session service on s's Session adapter. It runs the +// device, heartbeat and enrollment use cases. +func sessionService(t testing.TB, s *Store) *sessions.Service { + t.Helper() + service, err := sessions.NewService(sessionAdapter(s)) + if err != nil { + t.Fatal(err) + } + return service +} // stageTurnArtifacts stages export as the Turn's Artifacts through the Session // service on s's database, as the execution Worker does. diff --git a/services/core/internal/store/skill_version_deletion_test.go b/services/core/internal/store/skill_version_deletion_test.go index b609dfe8c..6f75c10d4 100644 --- a/services/core/internal/store/skill_version_deletion_test.go +++ b/services/core/internal/store/skill_version_deletion_test.go @@ -35,7 +35,7 @@ func TestSoleSkillVersionDeletionKeepsFrozenSetup(t *testing.T) { if err != nil { t.Fatal(err) } - frozen, err := s.ReadEnvironmentSetup(t.Context(), tenant, session.ID) + frozen, err := sessionAdapter(s).ReadEnvironmentSetup(t.Context(), tenant, session.ID) if err != nil || len(frozen.Skills) != 1 || frozen.Skills[0].Metadata.Version != "1" || !bytes.Equal(frozen.Skills[0].Archive, archive) { t.Fatal("fixture Session did not freeze the sole version", err) } @@ -43,7 +43,7 @@ func TestSoleSkillVersionDeletionKeepsFrozenSetup(t *testing.T) { t.Fatal(err) } - after, err := s.ReadEnvironmentSetup(t.Context(), tenant, session.ID) + after, err := sessionAdapter(s).ReadEnvironmentSetup(t.Context(), tenant, session.ID) if err != nil || !reflect.DeepEqual(after.Skills, frozen.Skills) { t.Fatal("frozen Session installation changed", err) } diff --git a/services/core/internal/store/structured_output_dispatch_test.go b/services/core/internal/store/structured_output_dispatch_test.go index 918f3ad4e..86bc24e64 100644 --- a/services/core/internal/store/structured_output_dispatch_test.go +++ b/services/core/internal/store/structured_output_dispatch_test.go @@ -28,7 +28,7 @@ func TestStructuredOutputDispatchRechecksOperationQualification(t *testing.T) { if err != nil { t.Fatal(err) } - if err = h.s.BindSessionDevice(t.Context(), h.tenant, h.session.ID, h.device.ID); err != nil { + if err = bindSessionDevice(t, h.db, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } caps := prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, StructuredOutput: proto.CapabilitySupported, MessageItems: proto.CapabilitySupported, Preparation: proto.CapabilitySupported}) diff --git a/services/core/internal/store/subagent_dispatch_test.go b/services/core/internal/store/subagent_dispatch_test.go index e49a26205..7b87bc5b8 100644 --- a/services/core/internal/store/subagent_dispatch_test.go +++ b/services/core/internal/store/subagent_dispatch_test.go @@ -24,7 +24,7 @@ func TestSubagentIdentityUsesLeasedDispatchJournal(t *testing.T) { if err != nil { t.Fatal(err) } - if err = h.s.BindSessionDevice(ctx, h.tenant, h.session.ID, h.device.ID); err != nil { + if err = bindSessionDevice(t, h.db, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } input := h.message("first", "root message") diff --git a/services/core/internal/store/subagent_identities_test.go b/services/core/internal/store/subagent_identities_test.go index 6b8a84bce..03857d54f 100644 --- a/services/core/internal/store/subagent_identities_test.go +++ b/services/core/internal/store/subagent_identities_test.go @@ -26,11 +26,11 @@ func TestSubagentIdentityIsAtomicScopedAndImmutable(t *testing.T) { journal := sessionExecution(t, w.lease) ctx := t.Context() tenant, session := newSubagentSession(t, s) - host, err := s.CreateDevice(ctx, tenant, "identity test", runtimedevice.HashCredential(uuid.NewString())) + host, err := sessionService(t, s).CreateDevice(ctx, tenant, "identity test", runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - if err = w.BindSessionDevice(ctx, tenant, session.ID, host.ID); err != nil { + if err = sessionExecution(t, w.lease).BindSessionDevice(ctx, tenant, session.ID, host.ID); err != nil { t.Fatal(err) } input := submitMessage(t, s, tenant, session.ID, "first") @@ -89,7 +89,7 @@ func TestSubagentIdentityIsAtomicScopedAndImmutable(t *testing.T) { if err != nil { t.Fatal(err) } - if err = w.BindSessionDevice(ctx, tenant, foreign.ID, host.ID); err != nil { + if err = sessionExecution(t, w.lease).BindSessionDevice(ctx, tenant, foreign.ID, host.ID); err != nil { t.Fatal(err) } foreignInput := submitMessage(t, s, tenant, foreign.ID, "first") diff --git a/services/core/internal/store/subagent_native_outputs_test.go b/services/core/internal/store/subagent_native_outputs_test.go index 51c3c92bd..caaf9946a 100644 --- a/services/core/internal/store/subagent_native_outputs_test.go +++ b/services/core/internal/store/subagent_native_outputs_test.go @@ -16,11 +16,11 @@ func TestSubagentNativeFunctionResultDoesNotConsumeOutputIndex(t *testing.T) { s, pool := testStore(t) owner := executionWriter(t, s) tenant, session := newSubagentSession(t, s) - host, err := s.CreateDevice(t.Context(), tenant, "child outputs", runtimedevice.HashCredential(uuid.NewString())) + host, err := sessionService(t, s).CreateDevice(t.Context(), tenant, "child outputs", runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - if err = owner.BindSessionDevice(t.Context(), tenant, session.ID, host.ID); err != nil { + if err = sessionExecution(t, owner.lease).BindSessionDevice(t.Context(), tenant, session.ID, host.ID); err != nil { t.Fatal(err) } input := submitMessage(t, s, tenant, session.ID, "start") @@ -90,11 +90,11 @@ func TestSubagentCancelledPartialMessageSurvivesHistoryReplay(t *testing.T) { s, _ := testStore(t) owner := executionWriter(t, s) tenant, session := newSubagentSession(t, s) - host, err := s.CreateDevice(t.Context(), tenant, "cancelled child", runtimedevice.HashCredential(uuid.NewString())) + host, err := sessionService(t, s).CreateDevice(t.Context(), tenant, "cancelled child", runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - if err = owner.BindSessionDevice(t.Context(), tenant, session.ID, host.ID); err != nil { + if err = sessionExecution(t, owner.lease).BindSessionDevice(t.Context(), tenant, session.ID, host.ID); err != nil { t.Fatal(err) } input := submitMessage(t, s, tenant, session.ID, "start") diff --git a/services/core/internal/store/subagent_resources_test.go b/services/core/internal/store/subagent_resources_test.go index d6e7a1655..49eecf66f 100644 --- a/services/core/internal/store/subagent_resources_test.go +++ b/services/core/internal/store/subagent_resources_test.go @@ -33,11 +33,11 @@ func TestSubagentResourcesNativeOwnershipLifecycleAndRecovery(t *testing.T) { journal := sessionExecution(t, owner.lease) ctx := t.Context() tenant, session := newSubagentSession(t, s) - host, err := s.CreateDevice(ctx, tenant, "child resources", runtimedevice.HashCredential(uuid.NewString())) + host, err := sessionService(t, s).CreateDevice(ctx, tenant, "child resources", runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - if err = owner.BindSessionDevice(ctx, tenant, session.ID, host.ID); err != nil { + if err = sessionExecution(t, owner.lease).BindSessionDevice(ctx, tenant, session.ID, host.ID); err != nil { t.Fatal(err) } root := submitMessage(t, s, tenant, session.ID, "first") diff --git a/services/core/internal/store/subagent_visibility_public_test.go b/services/core/internal/store/subagent_visibility_public_test.go index e7db1c169..80bc82abd 100644 --- a/services/core/internal/store/subagent_visibility_public_test.go +++ b/services/core/internal/store/subagent_visibility_public_test.go @@ -120,11 +120,11 @@ func TestSubagentVisibilityPublic(t *testing.T) { t.Fatal(page, err) } root := page.Turns[0].ID - host, err := s.CreateDevice(ctx, tenant, "subagent visibility", runtimedevice.HashCredential(uuid.NewString())) + host, err := fixtureSessionService(t, db).CreateDevice(ctx, tenant, "subagent visibility", runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - if err = writer.BindSessionDevice(ctx, tenant, session, host.ID); err != nil { + if err = leased.Sessions.BindSessionDevice(ctx, tenant, session, host.ID); err != nil { t.Fatal(err) } if _, err = writer.TransitionTurn(ctx, tenant, session, root, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { diff --git a/services/core/internal/store/template_composition_public_test.go b/services/core/internal/store/template_composition_public_test.go index 2741a894d..c4a697f2f 100644 --- a/services/core/internal/store/template_composition_public_test.go +++ b/services/core/internal/store/template_composition_public_test.go @@ -103,31 +103,31 @@ func TestTemplateCompositionOfficialClientPostgres(t *testing.T) { default: t.Fatalf("unknown case %q", label) } - for _, current := range []*store.Store{s, reopenedStore} { - setup, err := current.ReadEnvironmentSetup(t.Context(), tenant, id) - if err != nil || !reflect.DeepEqual(setup.Env, env) || !reflect.DeepEqual(setup.PackageMetadata(), packages) || len(setup.Commands) != len(commands) { - t.Fatalf("%s durable setup differs: %v", label, err) + // A reader built after the requests reads the frozen setup and files. + current := fixtureSessionStore(db) + setup, err := current.ReadEnvironmentSetup(t.Context(), tenant, id) + if err != nil || !reflect.DeepEqual(setup.Env, env) || !reflect.DeepEqual(setup.PackageMetadata(), packages) || len(setup.Commands) != len(commands) { + t.Fatalf("%s durable setup differs: %v", label, err) + } + for i, want := range commands { + if setup.Commands[i] != want { + t.Fatalf("%s command order differs", label) } - for i, want := range commands { - if setup.Commands[i] != want { - t.Fatalf("%s command order differs", label) - } + } + if _, err := current.ReadEnvironmentSetup(t.Context(), foreignTenant, id); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("%s foreign setup read: %v", label, err) + } + for position, want := range contents { + metadata, body, err := current.ReadInitialEnvironmentFile(t.Context(), tenant, id, position) + if err != nil || string(body) != want || metadata.Path != paths[position] || metadata.SizeBytes == nil || *metadata.SizeBytes != int64(len(want)) { + t.Fatalf("%s frozen file %d differs: %v", label, position, err) } - if _, err := current.ReadEnvironmentSetup(t.Context(), foreignTenant, id); !errors.Is(err, sessions.ErrNotFound) { - t.Fatalf("%s foreign setup read: %v", label, err) + if _, _, err := current.ReadInitialEnvironmentFile(t.Context(), foreignTenant, id, position); err == nil { + t.Fatalf("%s foreign file read succeeded", label) } - for position, want := range contents { - metadata, body, err := current.ReadInitialEnvironmentFile(t.Context(), tenant, id, position) - if err != nil || string(body) != want || metadata.Path != paths[position] || metadata.SizeBytes == nil || *metadata.SizeBytes != int64(len(want)) { - t.Fatalf("%s frozen file %d differs: %v", label, position, err) - } - if _, _, err := current.ReadInitialEnvironmentFile(t.Context(), foreignTenant, id, position); err == nil { - t.Fatalf("%s foreign file read succeeded", label) - } - var encrypted []byte - if err := pool.QueryRow(t.Context(), "SELECT contents FROM initial_environment_files WHERE id=$1", metadata.ID).Scan(&encrypted); err != nil || bytes.Contains(encrypted, []byte(marker)) { - t.Fatalf("%s plaintext file storage: %v", label, err) - } + var encrypted []byte + if err := pool.QueryRow(t.Context(), "SELECT contents FROM initial_environment_files WHERE id=$1", metadata.ID).Scan(&encrypted); err != nil || bytes.Contains(encrypted, []byte(marker)) { + t.Fatalf("%s plaintext file storage: %v", label, err) } } var encrypted, configuration []byte 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 5b9cb0c26..fb4a777ee 100644 --- a/services/core/internal/store/template_null_selection_public_test.go +++ b/services/core/internal/store/template_null_selection_public_test.go @@ -93,44 +93,44 @@ func TestTemplateNullSelectionOfficialClientPostgres(t *testing.T) { if !ok { t.Fatalf("missing expectation for %s", label) } - for _, current := range []*store.Store{s, reopenedStore} { - setup, err := current.ReadEnvironmentSetup(t.Context(), tenant, id) - if err != nil { - t.Fatalf("%s frozen setup: %v", label, err) - } - if !reflect.DeepEqual(setup.Env, map[string]string{"PRIVATE_SELECTION": marker}) || - !reflect.DeepEqual(setup.SkillMetadata(), want.Skills) || !reflect.DeepEqual(setup.PluginMetadata(), want.Plugins) || - !reflect.DeepEqual(append([]string{}, setup.CapabilityDirectories...), want.CapabilityDirectories) { - t.Fatalf("%s frozen selection changed", label) - } - if len(setup.Skills) != len(want.SkillDigests) || len(setup.Plugins) != len(want.PluginDigests) { - t.Fatalf("%s frozen archive count differs", label) - } - for i, skill := range setup.Skills { - digest := sha256.Sum256(skill.Archive) - if hex.EncodeToString(digest[:]) != want.SkillDigests[i] { - t.Fatalf("%s frozen Skill bytes changed", label) - } - } - for i, plugin := range setup.Plugins { - digest := sha256.Sum256(plugin.Archive) - if hex.EncodeToString(digest[:]) != want.PluginDigests[i] { - t.Fatalf("%s frozen Plugin bytes changed", label) - } - } - if _, err := current.ReadEnvironmentSetup(t.Context(), foreignTenant, id); !errors.Is(err, sessions.ErrNotFound) { - t.Fatalf("%s foreign setup read: %v", label, err) - } - file, body, err := current.ReadInitialEnvironmentFile(t.Context(), tenant, id, 0) - if err != nil || string(body) != marker+"-source" || file.Path != "/workspace/source.txt" { - t.Fatalf("%s frozen source file changed: %v", label, err) + // A reader built after the requests reads the frozen setup and files. + current := fixtureSessionStore(db) + setup, err := current.ReadEnvironmentSetup(t.Context(), tenant, id) + if err != nil { + t.Fatalf("%s frozen setup: %v", label, err) + } + if !reflect.DeepEqual(setup.Env, map[string]string{"PRIVATE_SELECTION": marker}) || + !reflect.DeepEqual(setup.SkillMetadata(), want.Skills) || !reflect.DeepEqual(setup.PluginMetadata(), want.Plugins) || + !reflect.DeepEqual(append([]string{}, setup.CapabilityDirectories...), want.CapabilityDirectories) { + t.Fatalf("%s frozen selection changed", label) + } + if len(setup.Skills) != len(want.SkillDigests) || len(setup.Plugins) != len(want.PluginDigests) { + t.Fatalf("%s frozen archive count differs", label) + } + for i, skill := range setup.Skills { + digest := sha256.Sum256(skill.Archive) + if hex.EncodeToString(digest[:]) != want.SkillDigests[i] { + t.Fatalf("%s frozen Skill bytes changed", label) } - if _, _, err := current.ReadInitialEnvironmentFile(t.Context(), foreignTenant, id, 0); err == nil { - t.Fatalf("%s foreign initial file read: %v", label, err) + } + for i, plugin := range setup.Plugins { + digest := sha256.Sum256(plugin.Archive) + if hex.EncodeToString(digest[:]) != want.PluginDigests[i] { + t.Fatalf("%s frozen Plugin bytes changed", label) } } + if _, err := current.ReadEnvironmentSetup(t.Context(), foreignTenant, id); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("%s foreign setup read: %v", label, err) + } + file, body, err := current.ReadInitialEnvironmentFile(t.Context(), tenant, id, 0) + if err != nil || string(body) != marker+"-source" || file.Path != "/workspace/source.txt" { + t.Fatalf("%s frozen source file changed: %v", label, err) + } + if _, _, err := current.ReadInitialEnvironmentFile(t.Context(), foreignTenant, id, 0); err == nil { + t.Fatalf("%s foreign initial file read: %v", label, err) + } var encryptedSetup, encryptedFile, configuration []byte - err := pool.QueryRow(t.Context(), "SELECT e.contents,f.contents,s.configuration FROM environment_setups e JOIN sessions s ON s.id=e.session_id JOIN initial_environment_files f ON f.session_id=s.id WHERE s.id=$1", id).Scan(&encryptedSetup, &encryptedFile, &configuration) + err = pool.QueryRow(t.Context(), "SELECT e.contents,f.contents,s.configuration FROM environment_setups e JOIN sessions s ON s.id=e.session_id JOIN initial_environment_files f ON f.session_id=s.id WHERE s.id=$1", id).Scan(&encryptedSetup, &encryptedFile, &configuration) if err != nil || len(encryptedSetup) == 0 || len(encryptedFile) == 0 || bytes.Contains(encryptedSetup, []byte(marker)) || bytes.Contains(encryptedFile, []byte(marker)) || bytes.Contains(configuration, []byte(marker)) { t.Fatalf("%s confidential initialization storage: %v", label, err) diff --git a/services/core/internal/store/turns_test.go b/services/core/internal/store/turns_test.go index 9b435ed49..297457606 100644 --- a/services/core/internal/store/turns_test.go +++ b/services/core/internal/store/turns_test.go @@ -4,10 +4,11 @@ import ( "context" "encoding/json" "errors" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "reflect" "sync" "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) func TestCancellationStaysBoundToItsOriginalTurn(t *testing.T) {