diff --git a/contracts/agents-api/core.openapi.yaml b/contracts/agents-api/core.openapi.yaml index 46e9e53ab..6699d80a6 100644 --- a/contracts/agents-api/core.openapi.yaml +++ b/contracts/agents-api/core.openapi.yaml @@ -4258,6 +4258,14 @@ paths: description: Conflict schema: $ref: '#/definitions/api.CoreErrorResponse' + "500": + description: Internal Server Error + schema: + $ref: '#/definitions/api.CoreErrorResponse' + "503": + description: Service Unavailable + schema: + $ref: '#/definitions/api.CoreErrorResponse' security: - DeploymentAdminAuth: [] summary: Get a self_hosted Session's installation commands diff --git a/contracts/agents-api/environment-executor-credentials.md b/contracts/agents-api/environment-executor-credentials.md index 5ace025ee..21cba485f 100644 --- a/contracts/agents-api/environment-executor-credentials.md +++ b/contracts/agents-api/environment-executor-credentials.md @@ -33,7 +33,7 @@ The installer calls these machine routes on Core: | `POST /api/v1/agent-daemon/installation` | Grant | The frozen binding: `version`, `protocol_version`, `environment_id`, `remote_url`, `workspace_directory`, `harness` | | `POST /api/v1/agent-daemon/installation/claim` | Grant | `{"executor_token":"SECRET"}`; 204 | -An invalid or expired grant returns 401 `installation_authorization_invalid`. Without matching installers the grant routes return 503 `installation_unavailable`. A malformed secret returns 400. Artifact routes carry no credential, and the grant is sent only to Core, never to an artifact host. +An invalid or expired grant returns 401 `installation_authorization_invalid`. Without matching installers the grant routes return 503 `installation_unavailable`. Core signs each grant with the installation's [`secrets/credential.key`](../../docs/configuration.md#installation-directory); without a configured key, the Session responses and Core-key read above and the grant routes return 503 `credential_storage_unavailable`. A malformed secret returns 400. Artifact routes carry no credential, and the grant is sent only to Core, never to an artifact host. ## Core-key routes diff --git a/contracts/agents-api/openapi.yaml b/contracts/agents-api/openapi.yaml index a7100a8de..64c942b2e 100644 --- a/contracts/agents-api/openapi.yaml +++ b/contracts/agents-api/openapi.yaml @@ -3270,6 +3270,10 @@ paths: description: Internal Server Error schema: $ref: '#/definitions/v1.ErrorResponse' + "503": + description: Service Unavailable + schema: + $ref: '#/definitions/v1.ErrorResponse' security: - BearerAuth: [] summary: Retrieve an execution Session @@ -3330,6 +3334,10 @@ paths: description: Internal Server Error schema: $ref: '#/definitions/v1.ErrorResponse' + "503": + description: Service Unavailable + schema: + $ref: '#/definitions/v1.ErrorResponse' security: - BearerAuth: [] summary: Update execution Session metadata diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 5253735bf..3d6aac68a 100644 --- a/services/core/IMPLEMENTATION.md +++ b/services/core/IMPLEMENTATION.md @@ -18,7 +18,7 @@ Shared vocabulary has one owner each, and domains use it rather than copy it. `i `internal/sessions` owns the Session vocabulary: Sessions, Turns, inputs, Environments and their provisioning failures, function calls, Item and Artifact reads, executor credentials, and the errors Session operations return, which `api` maps in `writeSessionsError`. It also owns the Session change vocabulary and its decisions: the public changes that report Turn and Session transitions, what a Turn that ends settles, measured Turn usage and the Session activity a change reports. Sequences that several Session operations share are `sessions` procedures, each over the transaction interface it declares: `CancelTurn` and `CancelWork` cancel work, `TrackInputActivity` reports the input activity a write changes, `FailEnvironment` and `TerminateEnvironment` settle an Environment that failed or whose managed compute ended, `CreateEnvironmentDevice` creates and binds a hosted Environment's dedicated Runtime device, `CheckComputeAdmission`, `CheckFileWriteGate` and `CheckInputStart` gate new work, `AdmitFunctionResult` saves a submitted function result for its Turn's call, and `LoadRequiredActions` lists the calls a Turn waits on. `AppendTurnEvents` records a batch of a Turn's execution observations, which `NewJournalBatch` validates, and `AppendTurnEvent` records its terminal outcome; both project the new entries. `ProjectSource` projects one journal entry or admitted input into Items, the root Turn's usage and Subagents, and `ProjectInput` projects an admitted input. `ValidTransition` decides the Turn status transitions, and the Subagent projection decides each Subagent's identity, lifecycle, child Turn and child Item changes. `sessions.ExecutionOperations` runs the Session writes only the execution owner makes; each function-call operation loads the Turn and its call in one leased Session transaction, decides in `sessions`, then applies in `sessionpg`, and `AppendTurnEvents` records and projects a Turn's journal batch in one leased Session transaction. `internal/items` owns public Items: it projects observations, merges them into stored Items and builds the ordered events that report each Item change. Neither imports persistence. `internal/persistence/postgres/sessionpg` loads the facts those decisions read and applies them inside the caller's Session transaction, under the Session lock: it allocates event sequence positions, event IDs, Item positions, output indexes and Subagent IDs, writes the change journal, each Turn's execution journal and its counters, Items, Turn usage, Subagent bindings, child Turns and child Items, and Artifact settlement, and prunes the change journal. It decides nothing. `WithSession` runs a Session operation's transaction on the pool or the lease: it locks the Session, applies the operation, then prunes the change journal. `BindSession` binds one Session to a transaction's queries as a `SessionTx`, which implements the procedures' interfaces. -Other adapters call only these `sessionpg` participants, on their own transaction's queries: `LockSession`, which orders a write against the Session's Turns; `BindSession`, which runs a Session procedure, such as `CreateEnvironmentDevice`, in that transaction; `LoadEnvironment` 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`. +Other adapters call only these `sessionpg` participants, on their own transaction's queries: `LockSession`, which orders a write against the Session's Turns; `BindSession`, which runs a Session procedure, such as `CreateEnvironmentDevice`, in that transaction; `LoadEnvironment`, which reads the tenant's Environment of a live Session; and `CreateSessionEnvironment`, which stores the Environment of a Session being created. The other `sessionpg` exports serve `store` until the Session operations leave it. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; deployment's placement fact-loading and reservation code sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. `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), 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`. +- `sessions` (`sessionpg`): Session use cases and reads as they leave `store`. The pooled `sessions.Service` runs the use cases on `sessionpg.Store`, which implements `sessions.Storage`, and plain reads use `sessions.Reader`, which `sessionpg.Store` also implements, directly. So far these cover root Items and Subagents (their reads), Session Artifacts (their reads, deletion and the staging of a Turn's export), Environments (their reads, the initialization list and the frozen setup and initial files, which `sessionpg.Store` opens with the credential key), devices (their reads, creation, revocation, heartbeats and Runtime enrollment), executor credentials (authentication and a Project's credential state through `sessions.ExecutorCredentialReader`, and issuance, rotation and revocation; the Core-key Project operations record their audit entry in the same transaction) and native installation authorization and claims, whose tokens `sessionpg.Store` signs with the credential key. `cmd/server` builds one `sessionpg.Store` with the credential key and wires the Service and the Store into `execution.Dispatcher.Sessions` and `SessionsReader`, into the api fields, into Runtime enrollment and into the daemon gateway; the Worker stages Artifacts through the Service; `cmd/environment-key` builds one without the key for its credential commands. Function calls and their application receipts, the Turn execution journal, Environment initialization, connection observations and their reconciliation, Session device binding, and file-write reservation and settlement run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. ## Request handling diff --git a/services/core/cmd/environment-key/main.go b/services/core/cmd/environment-key/main.go index 86b89fbfa..43c75db60 100644 --- a/services/core/cmd/environment-key/main.go +++ b/services/core/cmd/environment-key/main.go @@ -13,8 +13,9 @@ 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/identity" + "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/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" ) @@ -108,15 +109,19 @@ func run() error { return errors.New("invalid execution database configuration") } defer pool.Close() - s := store.New(pool) + // Executor credentials need no credential key. + credentials, err := sessions.NewService(sessionpg.New(pgunit.NewPool(pool), nil)) + if err != nil { + return err + } if options.revoke { - return credentialOperationError(s.RevokeExecutorCredential(ctx, options.principal, options.keyID)) + return credentialOperationError(credentials.RevokeExecutorCredential(ctx, options.principal, options.keyID)) } var credential sessions.IssuedExecutorCredential if options.rotate { - credential, err = s.RotateExecutorCredential(ctx, options.principal, options.keyID) + credential, err = credentials.RotateExecutorCredential(ctx, options.principal, options.keyID) } else { - credential, err = s.IssueExecutorCredential(ctx, options.principal, options.keyID, options.environment) + credential, err = credentials.IssueExecutorCredential(ctx, options.principal, options.keyID, options.environment) } if err != nil { return credentialOperationError(err) diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index 8c343b916..ce3c51a3a 100644 --- a/services/core/cmd/server/main.go +++ b/services/core/cmd/server/main.go @@ -392,7 +392,7 @@ func run() error { Artifacts: sessionService, ArtifactsReader: sessionStore, SessionAdmin: executionStore, - Environments: executionStore, EnvironmentsReader: sessionStore, ExecutorConnections: executorConnections{sessions: sessionStore, registry: registry}, + Environments: sessionService, EnvironmentsReader: sessionStore, ExecutorConnections: executorConnections{sessions: sessionStore, registry: registry}, Admin: executionStore, AdminAudit: auditStore, WriteAudit: auditStore, Metrics: metrics, RuntimeObservations: observationService, RuntimeHistory: historyService, } diff --git a/services/core/internal/api/environment_executor_management.go b/services/core/internal/api/environment_executor_management.go index 2b03da88a..9238a4c2e 100644 --- a/services/core/internal/api/environment_executor_management.go +++ b/services/core/internal/api/environment_executor_management.go @@ -66,9 +66,9 @@ func (h *Handler) listExecutorCredentials(w http.ResponseWriter, r *http.Request if !ok { return } - state, err := h.Environments.ProjectExecutorCredentialState(r.Context(), binding.Principal, chi.URLParam(r, "environment_id")) + state, err := h.EnvironmentsReader.ProjectExecutorCredentialState(r.Context(), binding.Principal, chi.URLParam(r, "environment_id")) if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } connection := ExecutorConnection{Status: "never_enrolled"} @@ -111,7 +111,7 @@ func (h *Handler) issueExecutorCredential(w http.ResponseWriter, r *http.Request var input EnvironmentExecutorCredentialRequest var fields map[string]json.RawMessage if decodeInputObject(raw, &input, "key_id", "rotate") != nil || json.Unmarshal(raw, &fields) != nil || bytes.Equal(bytes.TrimSpace(fields["rotate"]), []byte("null")) || !executorManagementID(input.KeyID) { - writeStoreError(w, r, sessions.ErrInvalidInput) + writeSessionsError(w, r, sessions.ErrInvalidInput) return } binding, ok := h.adminProjectScope(w, r) @@ -120,7 +120,7 @@ func (h *Handler) issueExecutorCredential(w http.ResponseWriter, r *http.Request } credential, err := h.Environments.IssueProjectExecutorCredential(r.Context(), binding.Principal, chi.URLParam(r, "environment_id"), input.KeyID, input.Rotate) if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } writeJSON(w, http.StatusCreated, credential) @@ -143,11 +143,11 @@ func (h *Handler) revokeExecutorCredential(w http.ResponseWriter, r *http.Reques } keyID := chi.URLParam(r, "key_id") if !executorManagementID(keyID) { - writeStoreError(w, r, sessions.ErrNotFound) + writeSessionsError(w, r, sessions.ErrNotFound) return } if err := h.Environments.RevokeProjectExecutorCredential(r.Context(), binding.Principal, chi.URLParam(r, "environment_id"), keyID); err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } w.Header().Set("Cache-Control", "no-store") diff --git a/services/core/internal/api/environment_executor_management_test.go b/services/core/internal/api/environment_executor_management_test.go index 25895cbc0..252a5bc83 100644 --- a/services/core/internal/api/environment_executor_management_test.go +++ b/services/core/internal/api/environment_executor_management_test.go @@ -54,7 +54,7 @@ func (f *executorManagementFixture) RevokeProjectExecutorCredential(ctx context. func executorManagementHandler(t *testing.T, key APIKey, f *executorManagementFixture, connected func(context.Context, string, string) (bool, error)) http.Handler { t.Helper() deps, fakes := managementFakes(t, key) - fakes.environments.projectExecutorCredentialState = f.ProjectExecutorCredentialState + fakes.environmentsReader.projectExecutorCredentialState = f.ProjectExecutorCredentialState fakes.environments.issueProjectExecutorCredential = f.IssueProjectExecutorCredential fakes.environments.revokeProjectExecutorCredential = f.RevokeProjectExecutorCredential fakes.executorConnections.executorConnected = connected diff --git a/services/core/internal/api/environment_installation.go b/services/core/internal/api/environment_installation.go index de20c3322..80412dcd8 100644 --- a/services/core/internal/api/environment_installation.go +++ b/services/core/internal/api/environment_installation.go @@ -90,12 +90,12 @@ func (h *Handler) installationAuthorization(w http.ResponseWriter, r *http.Reque } parts := strings.Fields(r.Header.Get("Authorization")) if len(r.Header.Values("Authorization")) != 1 || len(parts) != 2 || parts[0] != "Bearer" { - writeStoreError(w, r, sessions.ErrInstallationAuthorization) + writeSessionsError(w, r, sessions.ErrInstallationAuthorization) return sessions.InstallationAuthorization{}, "", false } claim, err := h.Environments.ValidateEnvironmentInstallation(r.Context(), parts[1], h.Execution.NativeInstaller.Version) if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return claim, "", false } return claim, parts[1], true @@ -115,7 +115,7 @@ func (h *Handler) prepareNativeInstallation(w http.ResponseWriter, r *http.Reque } environment, err := h.EnvironmentsReader.GetEnvironment(r.Context(), claim.Principal.TenantID, claim.Environment) if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } session, err := h.Sessions.GetSession(r.Context(), claim.Principal.TenantID, environment.SessionID) @@ -154,11 +154,11 @@ func (h *Handler) claimNativeInstallation(w http.ResponseWriter, r *http.Request } var input NativeInstallationClaim if decodeInputObject(raw, &input, "executor_token") != nil { - writeStoreError(w, r, sessions.ErrInvalidInput) + writeSessionsError(w, r, sessions.ErrInvalidInput) return } if err := h.Environments.ClaimEnvironmentInstallation(r.Context(), token, h.Execution.NativeInstaller.Version, input.ExecutorToken); err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } w.WriteHeader(http.StatusNoContent) @@ -171,7 +171,7 @@ func (h *Handler) claimNativeInstallation(w http.ResponseWriter, r *http.Request // @Param project_id path string true "Project UUID" // @Param environment_id path string true "Environment UUID" // @Success 200 {object} v1.EnvironmentInstallation -// @Failure 401,404,409 {object} CoreErrorResponse +// @Failure 401,404,409,500,503 {object} CoreErrorResponse // @Router /core/v1/projects/{project_id}/environments/{environment_id}/installation [get] func (h *Handler) getEnvironmentInstallation(w http.ResponseWriter, r *http.Request) { binding, ok := h.adminProjectScope(w, r) @@ -179,13 +179,13 @@ func (h *Handler) getEnvironmentInstallation(w http.ResponseWriter, r *http.Requ return } environment := chi.URLParam(r, "environment_id") - if _, err := h.Environments.ProjectExecutorCredentialState(r.Context(), binding.Principal, environment); err != nil { - writeStoreError(w, r, err) + if _, err := h.EnvironmentsReader.ProjectExecutorCredentialState(r.Context(), binding.Principal, environment); err != nil { + writeSessionsError(w, r, err) return } result, err := h.installationFor(r.Context(), binding.Principal, environment) if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } w.Header().Set("Cache-Control", "no-store") diff --git a/services/core/internal/api/environments.go b/services/core/internal/api/environments.go index d212919c9..0224efa06 100644 --- a/services/core/internal/api/environments.go +++ b/services/core/internal/api/environments.go @@ -12,9 +12,11 @@ import ( "github.com/go-chi/chi/v5" ) -// EnvironmentsReader reads Environments. +// EnvironmentsReader reads Environments and the executor credentials of a +// Project's self_hosted Environments. type EnvironmentsReader interface { GetEnvironment(context.Context, string, string) (sessions.Environment, error) + ProjectExecutorCredentialState(context.Context, identity.Principal, string) (sessions.ExecutorCredentialState, error) } // Environments grants and claims native installations, and manages the @@ -25,7 +27,6 @@ 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 - ProjectExecutorCredentialState(context.Context, identity.Principal, string) (sessions.ExecutorCredentialState, error) IssueProjectExecutorCredential(context.Context, identity.Principal, string, string, bool) (sessions.IssuedExecutorCredential, error) RevokeProjectExecutorCredential(context.Context, identity.Principal, string, string) error } diff --git a/services/core/internal/api/fakes_test.go b/services/core/internal/api/fakes_test.go index ac64473bd..a4fd14664 100644 --- a/services/core/internal/api/fakes_test.go +++ b/services/core/internal/api/fakes_test.go @@ -394,8 +394,9 @@ func (f *fakeEnvironmentWorkspaces) WriteEnvironmentFile(a0 context.Context, a1 } type fakeEnvironmentsReader struct { - t testing.TB - getEnvironment func(context.Context, string, string) (sessions.Environment, error) + t testing.TB + getEnvironment func(context.Context, string, string) (sessions.Environment, error) + projectExecutorCredentialState func(context.Context, identity.Principal, string) (sessions.ExecutorCredentialState, error) } func (f *fakeEnvironmentsReader) GetEnvironment(a0 context.Context, a1 string, a2 string) (sessions.Environment, error) { @@ -405,12 +406,18 @@ func (f *fakeEnvironmentsReader) GetEnvironment(a0 context.Context, a1 string, a return f.getEnvironment(a0, a1, a2) } +func (f *fakeEnvironmentsReader) ProjectExecutorCredentialState(a0 context.Context, a1 identity.Principal, a2 string) (sessions.ExecutorCredentialState, error) { + if f.projectExecutorCredentialState == nil { + unexpectedCall(f.t, "ProjectExecutorCredentialState") + } + return f.projectExecutorCredentialState(a0, a1, a2) +} + type fakeEnvironments struct { t testing.TB 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 - projectExecutorCredentialState func(context.Context, identity.Principal, string) (sessions.ExecutorCredentialState, error) issueProjectExecutorCredential func(context.Context, identity.Principal, string, string, bool) (sessions.IssuedExecutorCredential, error) revokeProjectExecutorCredential func(context.Context, identity.Principal, string, string) error } @@ -436,13 +443,6 @@ func (f *fakeEnvironments) ClaimEnvironmentInstallation(a0 context.Context, a1 s return f.claimEnvironmentInstallation(a0, a1, a2, a3) } -func (f *fakeEnvironments) ProjectExecutorCredentialState(a0 context.Context, a1 identity.Principal, a2 string) (sessions.ExecutorCredentialState, error) { - if f.projectExecutorCredentialState == nil { - unexpectedCall(f.t, "ProjectExecutorCredentialState") - } - return f.projectExecutorCredentialState(a0, a1, a2) -} - func (f *fakeEnvironments) IssueProjectExecutorCredential(a0 context.Context, a1 identity.Principal, a2 string, a3 string, a4 bool) (sessions.IssuedExecutorCredential, error) { if f.issueProjectExecutorCredential == nil { unexpectedCall(f.t, "IssueProjectExecutorCredential") diff --git a/services/core/internal/api/handler.go b/services/core/internal/api/handler.go index b135762b6..3892bbf83 100644 --- a/services/core/internal/api/handler.go +++ b/services/core/internal/api/handler.go @@ -298,7 +298,7 @@ func (h *Handler) createSession(w http.ResponseWriter, r *http.Request) { // @Param OpenAI-Beta header string true "agents=v1" // @Param session_id path string true "Session ID" // @Success 200 {object} v1.Session -// @Failure 400,401,404,500 {object} v1.ErrorResponse +// @Failure 400,401,404,500,503 {object} v1.ErrorResponse // @Router /agents/sessions/{session_id} [get] func (h *Handler) getSession(w http.ResponseWriter, r *http.Request) { session, err := h.Sessions.GetSession(r.Context(), tenantID(r), chi.URLParam(r, "session_id")) diff --git a/services/core/internal/api/session_metadata.go b/services/core/internal/api/session_metadata.go index 6daf53bac..e4c334fb5 100644 --- a/services/core/internal/api/session_metadata.go +++ b/services/core/internal/api/session_metadata.go @@ -23,7 +23,7 @@ import ( // @Param session_id path string true "Session ID" // @Param body body v1.UpdateSessionRequest true "Session metadata" // @Success 200 {object} v1.Session -// @Failure 400,401,404,413,500 {object} v1.ErrorResponse +// @Failure 400,401,404,413,500,503 {object} v1.ErrorResponse // @Router /agents/sessions/{session_id} [post] func (h *Handler) updateSession(w http.ResponseWriter, r *http.Request) { raw, ok := readJSONObject(w, r) diff --git a/services/core/internal/execution/environment_file_write.go b/services/core/internal/execution/environment_file_write.go index 89e3b379f..57d06f99e 100644 --- a/services/core/internal/execution/environment_file_write.go +++ b/services/core/internal/execution/environment_file_write.go @@ -109,7 +109,7 @@ func (w *Worker) runFileWrite(owner context.Context, request fileWriteRequest) f return fileWriteResult{err: sessions.ErrInvalidInput} } key := sessions.FileWriteIdentity{ID: uuid.NewString(), DeviceID: bound.ID, RequestSHA256: hex.EncodeToString(digest[:])} - intent, err := w.dispatcher.Store.ReserveEnvironmentFileWrite(ctx, environment.TenantID, environment.ID, key) + intent, err := w.dispatcher.sessionExecution.ReserveEnvironmentFileWrite(ctx, environment.TenantID, environment.ID, key) if err != nil { return fileWriteResult{err: err} } @@ -126,7 +126,7 @@ func (w *Worker) runFileWrite(owner context.Context, request fileWriteRequest) f } settle, stop := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) defer stop() - if _, err := w.dispatcher.Store.SettleEnvironmentFileWrite(settle, environment.TenantID, environment.ID, key, state); err != nil { + if _, err := w.dispatcher.sessionExecution.SettleEnvironmentFileWrite(settle, environment.TenantID, environment.ID, key, state); err != nil { return unavailable } if state == "rejected" { diff --git a/services/core/internal/persistence/postgres/sessionpg/environment.go b/services/core/internal/persistence/postgres/sessionpg/environment.go index bf290d46c..1374902eb 100644 --- a/services/core/internal/persistence/postgres/sessionpg/environment.go +++ b/services/core/internal/persistence/postgres/sessionpg/environment.go @@ -26,11 +26,10 @@ func LoadEnvironment(ctx context.Context, q *sqlc.Queries, tenant, environment s return EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, err) } -// LoadSessionEnvironment reads, on q, the Environment of the tenant's Session +// 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) { +// sessions.ErrInvalidInput. +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 @@ -56,7 +55,7 @@ func (s *Store) GetEnvironment(ctx context.Context, tenant, environment string) } func (s *Store) GetSessionEnvironment(ctx context.Context, tenant, session string) (sessions.Environment, error) { - return LoadSessionEnvironment(ctx, s.units.Queries(), tenant, session) + return loadSessionEnvironment(ctx, s.units.Queries(), tenant, session) } func (s *Store) ListEnvironmentInitializations(ctx context.Context, after string) ([]sessions.EnvironmentInitialization, error) { diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_file_writes.go b/services/core/internal/persistence/postgres/sessionpg/execution_file_writes.go new file mode 100644 index 000000000..14cbad562 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution_file_writes.go @@ -0,0 +1,176 @@ +package sessionpg + +import ( + "context" + "encoding/json" + "errors" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/writeaudit" +) + +// WithFileWriteReservation reads the Environment in the transaction before it +// locks the Environment's Session. +func (e *Execution) WithFileWriteReservation(ctx context.Context, tenant, environment string, apply func(context.Context, sessions.FileWriteReservationTx, sessions.Environment, sessions.LockedSession) error) error { + return e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + owner, err := LoadEnvironment(ctx, q, tenant, environment) + if err != nil { + return err + } + tenantID, err := parseID(tenant) + if err != nil { + return err + } + environmentID, err := parseID(owner.ID) + if err != nil { + return err + } + sessionID, err := parseID(owner.SessionID) + if err != nil { + return err + } + return inSession(ctx, q, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + return apply(ctx, newFileWriteTx(q, tenantID, sessionID, environmentID), owner, locked) + }) + }) +} + +// WithFileWriteSettlement reads the write in the transaction before it locks +// the Session of the write's Environment, whether or not that Session was +// publicly deleted. A malformed ID is sessions.ErrInvalidInput and an unknown +// write sessions.ErrNotFound. +func (e *Execution) WithFileWriteSettlement(ctx context.Context, tenant, environment, write string, apply func(context.Context, sessions.FileWriteSettlementTx) error) error { + tenantID, err := parseID(tenant) + if err != nil { + return err + } + environmentID, err := parseID(environment) + if err != nil { + return err + } + writeID, err := parseID(write) + if err != nil { + return err + } + return e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + row, err := q.GetEnvironmentFileWrite(ctx, sqlc.GetEnvironmentFileWriteParams{TenantID: tenantID, EnvironmentID: environmentID, ID: writeID}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.ErrNotFound + } + if err != nil { + return err + } + return inSession(ctx, q, tenantID, row.SessionID, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { + return apply(ctx, newFileWriteTx(q, tenantID, row.SessionID, environmentID)) + }) + }) +} + +// fileWriteTx is the file writes to the Session's Environment inside the +// Session transaction. +type fileWriteTx struct { + environmentTx + environment pgtype.UUID +} + +var ( + _ sessions.FileWriteReservationTx = (*fileWriteTx)(nil) + _ sessions.FileWriteSettlementTx = (*fileWriteTx)(nil) +) + +func newFileWriteTx(q *sqlc.Queries, tenant, session, environment pgtype.UUID) *fileWriteTx { + return &fileWriteTx{environmentTx: environmentTx{SessionTx: BindSession(q, tenant, session)}, environment: environment} +} + +func (t *fileWriteTx) LoadFileWrite(ctx context.Context, write string) (sessions.EnvironmentFileWrite, bool, error) { + row, err := t.loadFileWrite(ctx, write) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.EnvironmentFileWrite{}, false, nil + } + if err != nil { + return sessions.EnvironmentFileWrite{}, false, err + } + return fileWriteFromRow(row.EnvironmentFileWrite, row.SessionID), true, nil +} + +func (t *fileWriteTx) loadFileWrite(ctx context.Context, write string) (sqlc.GetEnvironmentFileWriteRow, error) { + id, err := parseID(write) + if err != nil { + return sqlc.GetEnvironmentFileWriteRow{}, err + } + return t.q.GetEnvironmentFileWrite(ctx, sqlc.GetEnvironmentFileWriteParams{TenantID: t.tenant, EnvironmentID: t.environment, ID: id}) +} + +func (t *fileWriteTx) LoadPendingInput(ctx context.Context) (bool, error) { + return t.q.EnvironmentFileWriteHasPendingInput(ctx, t.session) +} + +func (t *fileWriteTx) CreateFileWrite(ctx context.Context, key sessions.FileWriteIdentity) (sessions.EnvironmentFileWrite, error) { + id, err := parseID(key.ID) + if err != nil { + return sessions.EnvironmentFileWrite{}, err + } + device, err := parseID(key.DeviceID) + if err != nil { + return sessions.EnvironmentFileWrite{}, err + } + var origin []byte + if source, ok := writeaudit.FromContext(ctx); ok { + if origin, err = json.Marshal(source); err != nil { + return sessions.EnvironmentFileWrite{}, err + } + } + row, err := t.q.CreateEnvironmentFileWrite(ctx, sqlc.CreateEnvironmentFileWriteParams{ + ID: id, EnvironmentID: t.environment, DeviceID: device, RequestSha256: key.RequestSHA256, AuditSource: origin, + }) + if err != nil { + return sessions.EnvironmentFileWrite{}, err + } + return fileWriteFromRow(row, t.session), nil +} + +func (t *fileWriteTx) SettleFileWrite(ctx context.Context, write, state string) (sessions.EnvironmentFileWrite, error) { + id, err := parseID(write) + if err != nil { + return sessions.EnvironmentFileWrite{}, err + } + row, err := t.q.SettleEnvironmentFileWrite(ctx, sqlc.SettleEnvironmentFileWriteParams{EnvironmentID: t.environment, ID: id, State: state}) + if err != nil { + return sessions.EnvironmentFileWrite{}, err + } + return fileWriteFromRow(row, t.session), nil +} + +func (t *fileWriteTx) RecordFileWriteAudit(ctx context.Context, write string) error { + row, err := t.loadFileWrite(ctx, write) + if err != nil { + return err + } + if len(row.EnvironmentFileWrite.AuditSource) == 0 { + return nil + } + var source writeaudit.Source + if err := json.Unmarshal(row.EnvironmentFileWrite.AuditSource, &source); err != nil { + return err + } + return auditpg.RecordWriteAudit(writeaudit.WithSource(ctx, source), t.q, optionalID(t.tenant), "upload_file", "environment", optionalID(t.environment), optionalID(t.session)) +} + +func fileWriteFromRow(row sqlc.EnvironmentFileWrite, session pgtype.UUID) sessions.EnvironmentFileWrite { + result := sessions.EnvironmentFileWrite{ + Identity: sessions.FileWriteIdentity{ID: optionalID(row.ID), DeviceID: optionalID(row.DeviceID), RequestSHA256: row.RequestSha256}, + EnvironmentID: optionalID(row.EnvironmentID), SessionID: optionalID(session), State: row.State, CreatedAt: row.CreatedAt.Time, + } + if row.SettledAt.Valid { + settled := row.SettledAt.Time + result.SettledAt = &settled + } + return result +} diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_file_writes_test.go b/services/core/internal/persistence/postgres/sessionpg/execution_file_writes_test.go new file mode 100644 index 000000000..ece25f1f5 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution_file_writes_test.go @@ -0,0 +1,32 @@ +package sessionpg + +import ( + "errors" + "strings" + "testing" + + "github.com/google/uuid" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +func TestFileWriteSettlementNeedsAStoredWrite(t *testing.T) { + pool := pgtest.OpenIsolated(t, nil) + operations, _ := sessionExecution(t, pool) + tenantID, _, environmentID := newEnvironment(t, pool, "self_hosted", "connected") + tenant, environment := uuidText(tenantID), uuidText(environmentID) + key := sessions.FileWriteIdentity{ID: uuid.NewString(), DeviceID: uuid.NewString(), RequestSHA256: strings.Repeat("a", 64)} + for _, test := range []struct { + name, tenant, environment string + want error + }{ + {"an unknown write", tenant, environment, sessions.ErrNotFound}, + {"another tenant's Environment", uuid.NewString(), environment, sessions.ErrNotFound}, + {"a malformed Environment", tenant, "environment", sessions.ErrInvalidInput}, + } { + if _, err := operations.SettleEnvironmentFileWrite(t.Context(), test.tenant, test.environment, key, "committed"); !errors.Is(err, test.want) { + t.Fatalf("%s: %v, want %v", test.name, err, test.want) + } + } +} diff --git a/services/core/internal/persistence/postgres/sessionpg/executor_credentials.go b/services/core/internal/persistence/postgres/sessionpg/executor_credentials.go index 8238dd7e0..b47335ecd 100644 --- a/services/core/internal/persistence/postgres/sessionpg/executor_credentials.go +++ b/services/core/internal/persistence/postgres/sessionpg/executor_credentials.go @@ -2,6 +2,7 @@ package sessionpg import ( "context" + "crypto/hmac" "crypto/sha256" "encoding/hex" "errors" @@ -9,12 +10,20 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + "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/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/sessions" ) +// installationPurpose is the credential key purpose installation +// authorizations are signed under. +const installationPurpose = "environment-installation" + func (s *Store) AuthenticateEnvironmentExecutor(ctx context.Context, environment, digest string) (string, error) { id, err := pgunit.ParseID(environment) if err != nil { @@ -33,3 +42,285 @@ func (s *Store) AuthenticateEnvironmentExecutor(ctx context.Context, environment } return uuid.UUID(tenant.Bytes).String(), nil } + +func (s *Store) ProjectExecutorCredentialState(ctx context.Context, project identity.Principal, environment string) (sessions.ExecutorCredentialState, error) { + if err := project.Validate(); err != nil { + return sessions.ExecutorCredentialState{}, sessions.ErrInvalidInput + } + tenant, err := parseID(project.TenantID) + if err != nil { + return sessions.ExecutorCredentialState{}, err + } + environmentID := pgunit.PathID(environment) + result := sessions.ExecutorCredentialState{Credentials: []sessions.ExecutorCredential{}} + err = s.units.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + row, err := q.GetEnvironmentExecutorConnection(ctx, sqlc.GetEnvironmentExecutorConnectionParams{EnvironmentID: environmentID, TenantID: tenant}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.ErrNotFound + } + if err != nil { + return err + } + result.EnvironmentID = uuid.UUID(environmentID.Bytes).String() + result.Connection.EnvironmentStatus = row.EnvironmentStatus + result.Connection.DeviceID = optionalID(row.DeviceID) + if row.ExecutorKeyID.Valid { + key := optionalID(row.ExecutorKeyID) + result.Connection.BoundKeyID = &key + } + if row.EnrolledAt.Valid { + at := row.EnrolledAt.Time + result.Connection.EnrolledAt = &at + } + if row.LastSeenAt.Valid { + at := row.LastSeenAt.Time + result.Connection.LastSeenAt = &at + } + result.Connection.CredentialHash = row.CredentialHash.String + result.Credentials, err = listExecutorCredentials(ctx, q, tenant, environmentID, project.Subject()) + return err + }) + if err != nil { + return sessions.ExecutorCredentialState{}, err + } + return result, nil +} + +func listExecutorCredentials(ctx context.Context, q *sqlc.Queries, tenant, environment pgtype.UUID, subject identity.Subject) ([]sessions.ExecutorCredential, error) { + rows, err := q.ListEnvironmentExecutorCredentials(ctx, sqlc.ListEnvironmentExecutorCredentialsParams{ + TenantID: tenant, EnvironmentID: environment, SubjectKind: validText(subject.Kind), SubjectID: validText(subject.ID), + }) + if err != nil { + return nil, err + } + credentials := make([]sessions.ExecutorCredential, 0, len(rows)) + for _, row := range rows { + credential := sessions.ExecutorCredential{KeyID: optionalID(row.KeyID), CreatedAt: row.CreatedAt.Time} + if row.RevokedAt.Valid { + revoked := row.RevokedAt.Time + credential.RevokedAt = &revoked + } + credentials = append(credentials, credential) + } + return credentials, nil +} + +// validText is a present text value. +func validText(value string) pgtype.Text { return pgtype.Text{String: value, Valid: true} } + +func (s *Store) LoadProjectArchived(ctx context.Context, tenant string) (bool, error) { + id, err := parseID(tenant) + if err != nil { + return false, err + } + return loadProjectArchived(ctx, s.units.Queries(), id) +} + +// loadProjectArchived share-locks the tenant's Project, until the end of the +// transaction q is bound to, and reports whether it is archived. +func loadProjectArchived(ctx context.Context, q *sqlc.Queries, tenant pgtype.UUID) (bool, error) { + row, err := q.LockProjectByTenant(ctx, tenant) + if errors.Is(err, pgx.ErrNoRows) { + return false, sessions.ErrNotFound + } + if err != nil { + return false, err + } + return row.ArchivedAt.Valid, nil +} + +func (s *Store) ExecutorProjectExists(ctx context.Context, scope identity.ProjectScope) (bool, error) { + tenant, err := parseID(scope.TenantID) + if err != nil { + return false, err + } + return s.units.Queries().ExecutorProjectScopeExists(ctx, sqlc.ExecutorProjectScopeExistsParams{TenantID: tenant, OrganizationID: scope.OrganizationID, ProjectID: scope.ProjectID}) +} + +func (s *Store) LoadExecutorCredentialRestriction(ctx context.Context, principal identity.Principal, key string) (string, error) { + tenant, err := parseID(principal.TenantID) + if err != nil { + return "", err + } + id, err := parseID(key) + if err != nil { + return "", err + } + restriction, err := s.units.Queries().GetExecutorCredentialForPrincipal(ctx, sqlc.GetExecutorCredentialForPrincipalParams{ + KeyID: id, TenantID: tenant, SubjectKind: validText(principal.SubjectKind), SubjectID: validText(principal.SubjectID), + OrganizationID: principal.OrganizationID, ProjectID: principal.ProjectID, + }) + if errors.Is(err, pgx.ErrNoRows) { + return "", sessions.ErrNotFound + } + if err != nil { + return "", err + } + return optionalID(restriction), nil +} + +func (s *Store) SignInstallation(_ context.Context, payload string) (string, error) { + return s.installationSignature(payload) +} + +func (s *Store) VerifyInstallation(_ context.Context, payload, signature string) error { + want, err := s.installationSignature(payload) + if err != nil { + return err + } + if !hmac.Equal([]byte(want), []byte(signature)) { + return sessions.ErrInstallationAuthorization + } + return nil +} + +// installationSignature is the keyed digest of an installation authorization +// payload. A service without the credential key cannot sign or verify one. +func (s *Store) installationSignature(payload string) (string, error) { + if s.cipher == nil { + return "", credentialcrypto.ErrUnavailable + } + return s.cipher.Fingerprint(installationPurpose, payload) +} + +func (s *Store) WithExecutorCredentials(ctx context.Context, tenant string, apply func(context.Context, sessions.ExecutorCredentialTx) error) error { + tenantID, err := parseID(tenant) + if err != nil { + return err + } + return s.units.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + return apply(ctx, &executorCredentialTx{q: sqlc.New(tx), tenant: tenantID}) + }) +} + +// WithEnvironmentExecutorCredentials reads the Environment on the pool before +// it knows which Session to lock. +func (s *Store) WithEnvironmentExecutorCredentials(ctx context.Context, tenant, environment string, apply func(context.Context, sessions.EnvironmentExecutorCredentialTx, sessions.LockedSession) error) error { + tenantID, err := parseID(tenant) + if err != nil { + return err + } + if _, err := parseID(environment); err != nil { + return err + } + current, err := LoadEnvironment(ctx, s.units.Queries(), tenant, environment) + if err != nil { + return err + } + session := pgunit.PathID(current.SessionID) + return WithSession(ctx, s.units, tenantID, session, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + return apply(ctx, &environmentCredentialTx{executorCredentialTx: executorCredentialTx{q: q, tenant: tenantID}, session: session}, locked) + }) +} + +// executorCredentialTx is one tenant's executor credentials inside its +// caller's transaction. +type executorCredentialTx struct { + q *sqlc.Queries + tenant pgtype.UUID +} + +var _ sessions.ExecutorCredentialTx = (*executorCredentialTx)(nil) + +func (t *executorCredentialTx) LockProject(ctx context.Context) (bool, error) { + return loadProjectArchived(ctx, t.q, t.tenant) +} + +func (t *executorCredentialTx) ListExecutorCredentials(ctx context.Context, environment string, subject identity.Subject) ([]sessions.ExecutorCredential, error) { + id, err := parseID(environment) + if err != nil { + return nil, err + } + return listExecutorCredentials(ctx, t.q, t.tenant, id, subject) +} + +func (t *executorCredentialTx) AuthenticateExecutor(ctx context.Context, environment, digest string) (bool, error) { + id, err := parseID(environment) + if err != nil { + return false, err + } + _, err = t.q.AuthenticateEnvironmentExecutor(ctx, sqlc.AuthenticateEnvironmentExecutorParams{EnvironmentID: id, TokenSha256: digest}) + if errors.Is(err, pgx.ErrNoRows) { + return false, nil + } + return err == nil, err +} + +func (t *executorCredentialTx) IssueExecutorCredential(ctx context.Context, grant sessions.ExecutorCredentialGrant) (sessions.IssuedExecutorCredential, error) { + key, err := parseID(grant.KeyID) + if err != nil { + return sessions.IssuedExecutorCredential{}, err + } + var restriction pgtype.UUID + if grant.EnvironmentID != "" { + if restriction, err = parseID(grant.EnvironmentID); err != nil { + return sessions.IssuedExecutorCredential{}, err + } + } + row, err := t.q.IssueExecutorCredential(ctx, sqlc.IssueExecutorCredentialParams{ + KeyID: key, TenantID: t.tenant, SubjectKind: validText(grant.Principal.SubjectKind), SubjectID: validText(grant.Principal.SubjectID), + OrganizationID: grant.Principal.OrganizationID, ProjectID: grant.Principal.ProjectID, EnvironmentID: restriction, TokenSha256: grant.Digest, + }) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.IssuedExecutorCredential{}, sessions.ErrExecutorCredentialExists + } + if err != nil { + return sessions.IssuedExecutorCredential{}, err + } + return sessions.IssuedExecutorCredential{KeyID: optionalID(row.KeyID), EnvironmentID: optionalID(row.EnvironmentID)}, nil +} + +func (t *executorCredentialTx) RotateExecutorCredential(ctx context.Context, subject identity.Subject, key, digest string) (sessions.IssuedExecutorCredential, error) { + id, err := parseID(key) + if err != nil { + return sessions.IssuedExecutorCredential{}, err + } + row, err := t.q.RotateExecutorCredential(ctx, sqlc.RotateExecutorCredentialParams{KeyID: id, TenantID: t.tenant, SubjectKind: validText(subject.Kind), SubjectID: validText(subject.ID), TokenSha256: digest}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.IssuedExecutorCredential{}, sessions.ErrNotFound + } + if err != nil { + return sessions.IssuedExecutorCredential{}, err + } + return sessions.IssuedExecutorCredential{KeyID: optionalID(row.KeyID), EnvironmentID: optionalID(row.EnvironmentID)}, nil +} + +func (t *executorCredentialTx) RevokeExecutorCredential(ctx context.Context, subject identity.Subject, key string) error { + id, err := parseID(key) + if err != nil { + return err + } + n, err := t.q.RevokeExecutorCredential(ctx, sqlc.RevokeExecutorCredentialParams{KeyID: id, TenantID: t.tenant, SubjectKind: validText(subject.Kind), SubjectID: validText(subject.ID)}) + if err == nil && n == 0 { + return sessions.ErrNotFound + } + return err +} + +func (t *executorCredentialTx) RecordExecutorCredentialAudit(ctx context.Context, action, key string) error { + return auditpg.RecordAdminMutation(ctx, t.q, optionalID(t.tenant), action, "executor_credential", key) +} + +// environmentCredentialTx is the executor credentials of the Session's +// tenant inside the Session's transaction. +type environmentCredentialTx struct { + executorCredentialTx + session pgtype.UUID +} + +var _ sessions.EnvironmentExecutorCredentialTx = (*environmentCredentialTx)(nil) + +func (t *environmentCredentialTx) LoadSessionCreator(ctx context.Context) (identity.Subject, bool, error) { + row, err := t.q.GetSession(ctx, sqlc.GetSessionParams{TenantID: t.tenant, ID: t.session}) + if errors.Is(err, pgx.ErrNoRows) { + return identity.Subject{}, false, sessions.ErrNotFound + } + if err != nil { + return identity.Subject{}, false, err + } + if !row.CreatorKind.Valid || !row.CreatorID.Valid { + return identity.Subject{}, false, nil + } + return identity.Subject{Kind: row.CreatorKind.String, ID: row.CreatorID.String}, true, nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/executor_credentials_test.go b/services/core/internal/persistence/postgres/sessionpg/executor_credentials_test.go new file mode 100644 index 000000000..fc5f3ade4 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/executor_credentials_test.go @@ -0,0 +1,93 @@ +package sessionpg + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + + "github.com/google/uuid" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" + "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 TestInstallationSignaturesUseTheCredentialKey(t *testing.T) { + keyed := New(nil, frozenCipher(t)) + signature, err := keyed.SignInstallation(t.Context(), "payload") + if err != nil { + t.Fatal(err) + } + if err := keyed.VerifyInstallation(t.Context(), "payload", signature); err != nil { + t.Fatal(err) + } + if err := keyed.VerifyInstallation(t.Context(), "other", signature); !errors.Is(err, sessions.ErrInstallationAuthorization) { + t.Fatalf("another payload: %v", err) + } + keyless := New(nil, nil) + if _, err := keyless.SignInstallation(t.Context(), "payload"); !errors.Is(err, credentialcrypto.ErrUnavailable) { + t.Fatalf("keyless signature: %v", err) + } + if err := keyless.VerifyInstallation(t.Context(), "payload", signature); !errors.Is(err, credentialcrypto.ErrUnavailable) { + t.Fatalf("keyless verification: %v", err) + } +} + +func TestExecutorCredentialTxTranslatesOutcomes(t *testing.T) { + pool := pgtest.Open(t) + store := New(pgunit.NewPool(pool), nil) + tenantID, sessionID, environmentID := newEnvironment(t, pool, "self_hosted", "pending") + tenant, environment := uuidText(tenantID), uuidText(environmentID) + exec(t, pool, `INSERT INTO execution_project_scopes(tenant_id, organization_id, project_id) VALUES ($1, 'org', $2)`, tenantID, tenant) + principal := identity.Principal{ProjectScope: identity.ProjectScope{TenantID: tenant, OrganizationID: "org", ProjectID: tenant}, SubjectKind: "user", SubjectID: "owner"} + within := func(apply func(context.Context, sessions.EnvironmentExecutorCredentialTx) error) { + t.Helper() + err := store.WithEnvironmentExecutorCredentials(t.Context(), tenant, environment, func(ctx context.Context, tx sessions.EnvironmentExecutorCredentialTx, _ sessions.LockedSession) error { + return apply(ctx, tx) + }) + if err != nil { + t.Fatal(err) + } + } + within(func(ctx context.Context, tx sessions.EnvironmentExecutorCredentialTx) error { + if _, found, err := tx.LoadSessionCreator(ctx); err != nil || found { + return fmt.Errorf("creator of a Session without one: %v %v", found, err) + } + return nil + }) + exec(t, pool, `UPDATE sessions SET creator_kind = 'user', creator_id = 'owner' WHERE id = $1`, sessionID) + // Digests are unique across tenants, so each run presents its own. + key, digest := uuid.NewString(), strings.ReplaceAll(uuid.NewString()+uuid.NewString(), "-", "") + other := identity.Subject{Kind: "user", ID: "other"} + within(func(ctx context.Context, tx sessions.EnvironmentExecutorCredentialTx) error { + if creator, found, err := tx.LoadSessionCreator(ctx); err != nil || !found || creator != principal.Subject() { + return fmt.Errorf("creator %+v %v %v", creator, found, err) + } + if _, err := tx.LockProject(ctx); !errors.Is(err, sessions.ErrNotFound) { + return fmt.Errorf("lock of a missing Project: %v", err) + } + grant := sessions.ExecutorCredentialGrant{Principal: principal, KeyID: key, EnvironmentID: environment, Digest: digest} + if issued, err := tx.IssueExecutorCredential(ctx, grant); err != nil || issued != (sessions.IssuedExecutorCredential{KeyID: key, EnvironmentID: environment}) { + return fmt.Errorf("issue %+v %v", issued, err) + } + if _, err := tx.IssueExecutorCredential(ctx, grant); !errors.Is(err, sessions.ErrExecutorCredentialExists) { + return fmt.Errorf("reissue: %v", err) + } + for presented, want := range map[string]bool{digest: true, strings.Repeat("b", 64): false} { + if current, err := tx.AuthenticateExecutor(ctx, environment, presented); err != nil || current != want { + return fmt.Errorf("authenticate %s: %v %v", presented, current, err) + } + } + if _, err := tx.RotateExecutorCredential(ctx, other, key, digest); !errors.Is(err, sessions.ErrNotFound) { + return fmt.Errorf("another subject's rotation: %v", err) + } + if err := tx.RevokeExecutorCredential(ctx, other, key); !errors.Is(err, sessions.ErrNotFound) { + return fmt.Errorf("another subject's revocation: %v", err) + } + return nil + }) +} diff --git a/services/core/internal/persistence/postgres/sessionpg/store.go b/services/core/internal/persistence/postgres/sessionpg/store.go index f595a647f..0ef7d09b7 100644 --- a/services/core/internal/persistence/postgres/sessionpg/store.go +++ b/services/core/internal/persistence/postgres/sessionpg/store.go @@ -7,15 +7,16 @@ import ( ) // Store is the pooled Session adapter: the storage of the Session use cases -// and the Session reads. cipher opens the Session data frozen at creation; it -// is nil on a service without a credential key. +// and the Session reads. cipher opens the Session data frozen at creation and +// signs installation authorizations; 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, opening frozen Session data -// with cipher. +// and signing installation authorizations with cipher. func New(units *pgunit.Pool, cipher *credentialcrypto.Cipher) *Store { return &Store{units: units, cipher: cipher} } diff --git a/services/core/internal/sandbox/providers/configuration_flow_test.go b/services/core/internal/sandbox/providers/configuration_flow_test.go index 88d4f46d3..27ad2b0ee 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, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, Admin: s, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, + Environments: struct{ api.Environments }{}, 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 0b2b4561c..7132974d7 100644 --- a/services/core/internal/sessions/artifacts_test.go +++ b/services/core/internal/sessions/artifacts_test.go @@ -10,6 +10,8 @@ import ( "strings" "testing" "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" ) // exportFile is one entry of a test export. @@ -215,6 +217,46 @@ func (s *fakeArtifactStorage) WithEnrollment(context.Context, string, string, fu return nil } +func (s *fakeArtifactStorage) GetEnvironment(context.Context, string, string) (Environment, error) { + s.t.Fatal("unexpected call to GetEnvironment") + return Environment{}, nil +} + +func (s *fakeArtifactStorage) LoadProjectArchived(context.Context, string) (bool, error) { + s.t.Fatal("unexpected call to LoadProjectArchived") + return false, nil +} + +func (s *fakeArtifactStorage) ExecutorProjectExists(context.Context, identity.ProjectScope) (bool, error) { + s.t.Fatal("unexpected call to ExecutorProjectExists") + return false, nil +} + +func (s *fakeArtifactStorage) LoadExecutorCredentialRestriction(context.Context, identity.Principal, string) (string, error) { + s.t.Fatal("unexpected call to LoadExecutorCredentialRestriction") + return "", nil +} + +func (s *fakeArtifactStorage) SignInstallation(context.Context, string) (string, error) { + s.t.Fatal("unexpected call to SignInstallation") + return "", nil +} + +func (s *fakeArtifactStorage) VerifyInstallation(context.Context, string, string) error { + s.t.Fatal("unexpected call to VerifyInstallation") + return nil +} + +func (s *fakeArtifactStorage) WithExecutorCredentials(context.Context, string, func(context.Context, ExecutorCredentialTx) error) error { + s.t.Fatal("unexpected call to WithExecutorCredentials") + return nil +} + +func (s *fakeArtifactStorage) WithEnvironmentExecutorCredentials(context.Context, string, string, func(context.Context, EnvironmentExecutorCredentialTx, LockedSession) error) error { + s.t.Fatal("unexpected call to WithEnvironmentExecutorCredentials") + 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_test.go b/services/core/internal/sessions/devices_test.go index 00b8fdd64..20f8676e5 100644 --- a/services/core/internal/sessions/devices_test.go +++ b/services/core/internal/sessions/devices_test.go @@ -9,10 +9,11 @@ import ( "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 { +// fakeStorage is strict pooled Session storage for the device and executor +// credential 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 fakeStorage struct { t *testing.T calls []string @@ -25,9 +26,23 @@ type fakeDeviceStorage struct { enrollment *fakeTx environment Environment locked LockedSession + + getEnvironment func() (Environment, error) + loadProjectArchived func() (bool, error) + projectExists func() (bool, error) + loadRestriction func() (string, error) + // installationKey signs installation authorizations; empty fails the + // test. + installationKey string + // verifyInstallation, when set, is what VerifyInstallation returns in + // place of checking the signature. + verifyInstallation error + // credentials is the transaction the executor credential With methods + // apply in, with locked. + credentials *fakeTx } -func (s *fakeDeviceStorage) record(name string, set bool, detail ...string) { +func (s *fakeStorage) record(name string, set bool, detail ...string) { s.t.Helper() if !set { s.t.Fatalf("unexpected call to %s", name) @@ -35,42 +50,42 @@ func (s *fakeDeviceStorage) record(name string, set bool, detail ...string) { s.calls = append(s.calls, strings.Join(append([]string{name}, detail...), " ")) } -func (s *fakeDeviceStorage) CreateDevice(_ context.Context, tenant string, registration DeviceRegistration) (ExecutionDevice, error) { +func (s *fakeStorage) 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 { +func (s *fakeStorage) 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) { +func (s *fakeStorage) 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) { +func (s *fakeStorage) 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 { +func (s *fakeStorage) 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 { +func (s *fakeStorage) 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 { +func (s *fakeStorage) DeleteSessionArtifact(context.Context, string, string, string) error { s.t.Fatal("unexpected call to DeleteSessionArtifact") return nil } -func deviceService(t *testing.T, storage *fakeDeviceStorage) *Service { +func deviceService(t *testing.T, storage *fakeStorage) *Service { t.Helper() service, err := NewService(storage) if err != nil { @@ -106,7 +121,7 @@ func TestNewDeviceRegistration(t *testing.T) { } func TestDeviceUseCases(t *testing.T) { - storage := &fakeDeviceStorage{t: t} + storage := &fakeStorage{t: t} if _, err := deviceService(t, storage).CreateDevice(t.Context(), "tenant", "runtime", "secret"); !errors.Is(err, ErrInvalidInput) { t.Fatalf("invalid registration: %v", err) } @@ -127,7 +142,7 @@ func TestDeviceUseCases(t *testing.T) { // 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)} + storage := &fakeStorage{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}) { @@ -142,7 +157,7 @@ func TestDeviceUseCases(t *testing.T) { } } failing := func() (bool, error) { return true, errStorage } - storage = &fakeDeviceStorage{t: t, touchDevice: failing, touchAuthenticatedDevice: failing} + storage = &fakeStorage{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) } @@ -174,7 +189,7 @@ func TestEnrollRuntimeBindsUnderTheSessionLock(t *testing.T) { 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} + storage := &fakeStorage{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) diff --git a/services/core/internal/sessions/environment_installation.go b/services/core/internal/sessions/environment_installation.go new file mode 100644 index 000000000..c0809381f --- /dev/null +++ b/services/core/internal/sessions/environment_installation.go @@ -0,0 +1,130 @@ +package sessions + +import ( + "context" + "encoding/base64" + "encoding/json" + "strings" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" +) + +// InstallationAuthorization permits claiming one Environment's connect-only key. +// The Environment UUID is reserved as that key's ID. Reissuing an authorization +// never rotates or revives the key, and a retry must prove the same local secret. +type InstallationAuthorization struct { + Principal identity.Principal `json:"principal"` + Environment string `json:"environment_id"` + Version string `json:"version"` + ExpiresAt int64 `json:"expires_at"` +} + +const ( + // installationLifetime is how long an installation authorization lasts. + installationLifetime = 30 * time.Minute + // maxInstallationToken bounds the installation token a request presents. + maxInstallationToken = 4096 +) + +// AuthorizeEnvironmentInstallation authorizes the installer of version to +// claim the connect-only key of the principal's self_hosted Environment, and +// returns the token and when it expires, in Unix seconds. The token is the +// base64url JSON authorization and its signature, joined by a dot. An archived +// Project gets none (projects.ErrArchived). +func (s *Service) AuthorizeEnvironmentInstallation(ctx context.Context, principal identity.Principal, environment, version string) (string, int64, error) { + if _, err := s.selfHostedTarget(ctx, principal, environment); err != nil { + return "", 0, err + } + if err := s.checkActiveProject(ctx, principal.TenantID); err != nil { + return "", 0, err + } + claim := InstallationAuthorization{Principal: principal, Environment: environment, Version: version, ExpiresAt: time.Now().Add(installationLifetime).Unix()} + payload, err := json.Marshal(claim) + if err != nil { + return "", 0, err + } + encoded := base64.RawURLEncoding.EncodeToString(payload) + signature, err := s.storage.SignInstallation(ctx, encoded) + if err != nil { + return "", 0, err + } + return encoded + "." + signature, claim.ExpiresAt, nil +} + +// ValidateEnvironmentInstallation returns the authorization the token carries +// while it is signed, unexpired, for version, and still names a self_hosted +// Environment of an active Project. Any other token is +// ErrInstallationAuthorization. +func (s *Service) ValidateEnvironmentInstallation(ctx context.Context, token, version string) (InstallationAuthorization, error) { + claim, _, err := s.validateInstallation(ctx, token, version) + return claim, err +} + +// validateInstallation also returns the Environment the authorization names. +func (s *Service) validateInstallation(ctx context.Context, token, version string) (InstallationAuthorization, Environment, error) { + encoded, signature, ok := strings.Cut(token, ".") + if !ok || len(token) > maxInstallationToken { + return InstallationAuthorization{}, Environment{}, ErrInstallationAuthorization + } + if err := s.storage.VerifyInstallation(ctx, encoded, signature); err != nil { + return InstallationAuthorization{}, Environment{}, err + } + var claim InstallationAuthorization + payload, err := base64.RawURLEncoding.DecodeString(encoded) + if err != nil || json.Unmarshal(payload, &claim) != nil || claim.Version != version || claim.ExpiresAt <= time.Now().Unix() { + return InstallationAuthorization{}, Environment{}, ErrInstallationAuthorization + } + target, err := s.selfHostedTarget(ctx, claim.Principal, claim.Environment) + if err != nil { + return InstallationAuthorization{}, Environment{}, ErrInstallationAuthorization + } + if err := s.checkActiveProject(ctx, claim.Principal.TenantID); err != nil { + return InstallationAuthorization{}, Environment{}, ErrInstallationAuthorization + } + return claim, target, nil +} + +// ClaimEnvironmentInstallation stores only the digest of a secret the +// installer generated and persisted before this request. A lost response is +// safe to retry with the same secret; another machine, or a revoked or rotated +// key, cannot claim it again (ErrExecutorCredentialExists). +func (s *Service) ClaimEnvironmentInstallation(ctx context.Context, token, version, secret string) error { + claim, target, err := s.validateInstallation(ctx, token, version) + if err != nil { + return err + } + digest, err := installationSecretDigest(secret) + if err != nil { + return err + } + principal := claim.Principal + // The Environment's ID is reserved as the ID of the key its installation + // claims. + key := target.ID + return s.withExecutorTarget(ctx, principal, target.ID, func(ctx context.Context, tx ExecutorCredentialTx) error { + if err := lockActiveProject(ctx, tx); err != nil { + return err + } + keys, err := tx.ListExecutorCredentials(ctx, target.ID, principal.Subject()) + if err != nil { + return err + } + if len(keys) > 0 { + // Onboarding never replaces an operator-issued key. + if len(keys) != 1 || keys[0].KeyID != key || keys[0].RevokedAt != nil { + return ErrExecutorCredentialExists + } + current, err := tx.AuthenticateExecutor(ctx, target.ID, digest) + if err != nil { + return err + } + if !current { + return ErrExecutorCredentialExists + } + return nil + } + _, err = tx.IssueExecutorCredential(ctx, ExecutorCredentialGrant{Principal: principal, KeyID: key, EnvironmentID: target.ID, Digest: digest}) + return err + }) +} diff --git a/services/core/internal/sessions/environment_installation_test.go b/services/core/internal/sessions/environment_installation_test.go new file mode 100644 index 000000000..6a2a3b1a4 --- /dev/null +++ b/services/core/internal/sessions/environment_installation_test.go @@ -0,0 +1,164 @@ +package sessions + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "errors" + "strings" + "testing" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" +) + +// installationToken signs claim as storage does. +func installationToken(t *testing.T, storage *fakeStorage, claim InstallationAuthorization) string { + t.Helper() + payload, err := json.Marshal(claim) + if err != nil { + t.Fatal(err) + } + encoded := base64.RawURLEncoding.EncodeToString(payload) + return encoded + "." + storage.signature(encoded) +} + +func assertStorageCalls(t *testing.T, storage *fakeStorage, want ...string) { + t.Helper() + if strings.Join(storage.calls, "\n") != strings.Join(want, "\n") { + t.Fatalf("storage calls:\n%s\nwant:\n%s", strings.Join(storage.calls, "\n"), strings.Join(want, "\n")) + } +} + +func TestEnvironmentInstallationAuthorization(t *testing.T) { + target := "GetEnvironment " + testTenant + " " + testEnvironment + archive := "LoadProjectArchived " + testTenant + installable := func() *fakeStorage { + return &fakeStorage{t: t, installationKey: "key", getEnvironment: returns(selfHosted), loadProjectArchived: returns(false)} + } + storage := installable() + service := deviceService(t, storage) + token, expires, err := service.AuthorizeEnvironmentInstallation(t.Context(), testPrincipal, testEnvironment, "1.2.3") + if err != nil { + t.Fatal(err) + } + if lifetime := time.Until(time.Unix(expires, 0)); lifetime <= 29*time.Minute || lifetime > 30*time.Minute { + t.Fatalf("expires in %s", lifetime) + } + claim, err := service.ValidateEnvironmentInstallation(t.Context(), token, "1.2.3") + if want := (InstallationAuthorization{Principal: testPrincipal, Environment: testEnvironment, Version: "1.2.3", ExpiresAt: expires}); err != nil || claim != want { + t.Fatalf("claim %+v, %v", claim, err) + } + assertStorageCalls(t, storage, target, archive, "SignInstallation", "VerifyInstallation", target, archive) + + encoded, signature, _ := strings.Cut(token, ".") + expired := installationToken(t, storage, InstallationAuthorization{Principal: testPrincipal, Environment: testEnvironment, Version: "1.2.3", ExpiresAt: time.Now().Add(-time.Second).Unix()}) + for _, test := range []struct { + name, token, version string + storage *fakeStorage + }{ + {"another version", token, "1.2.4", installable()}, + {"another signature", encoded + "." + strings.Repeat("0", len(signature)), "1.2.3", installable()}, + {"no signature", encoded, "1.2.3", installable()}, + {"too long", token + strings.Repeat("0", maxInstallationToken), "1.2.3", installable()}, + {"expired", expired, "1.2.3", installable()}, + {"deleted Environment", token, "1.2.3", &fakeStorage{t: t, installationKey: "key", getEnvironment: fails[Environment](ErrNotFound)}}, + {"archived Project", token, "1.2.3", &fakeStorage{t: t, installationKey: "key", getEnvironment: returns(selfHosted), loadProjectArchived: returns(true)}}, + } { + if _, err := deviceService(t, test.storage).ValidateEnvironmentInstallation(t.Context(), test.token, test.version); !errors.Is(err, ErrInstallationAuthorization) { + t.Fatalf("%s: %v", test.name, err) + } + } + + hosted := &fakeStorage{t: t, getEnvironment: returns(hostedEnvironment)} + if _, _, err := deviceService(t, hosted).AuthorizeEnvironmentInstallation(t.Context(), testPrincipal, testEnvironment, "1.2.3"); !errors.Is(err, ErrNotFound) { + t.Fatalf("OpenAI-hosted Environment: %v", err) + } + archived := &fakeStorage{t: t, getEnvironment: returns(selfHosted), loadProjectArchived: returns(true)} + if _, _, err := deviceService(t, archived).AuthorizeEnvironmentInstallation(t.Context(), testPrincipal, testEnvironment, "1.2.3"); !errors.Is(err, projects.ErrArchived) { + t.Fatalf("archived Project: %v", err) + } +} + +func TestClaimEnvironmentInstallation(t *testing.T) { + secret := base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{7}, 32)) + digest := executorDigest(secret) + validation := []string{"VerifyInstallation", "GetEnvironment " + testTenant + " " + testEnvironment, "LoadProjectArchived " + testTenant} + claiming := append(validation[:len(validation):len(validation)], "WithEnvironmentExecutorCredentials "+testTenant+" "+testEnvironment) + checked := []string{"LoadSessionCreator", "LockProject", "ListExecutorCredentials " + testEnvironment + " creator"} + claimed := func(grant ExecutorCredentialGrant) (IssuedExecutorCredential, error) { + if grant != (ExecutorCredentialGrant{Principal: testPrincipal, KeyID: testEnvironment, EnvironmentID: testEnvironment, Digest: digest}) { + t.Errorf("grant %+v", grant) + } + return IssuedExecutorCredential{KeyID: grant.KeyID, EnvironmentID: grant.EnvironmentID}, nil + } + own := []ExecutorCredential{{KeyID: testEnvironment}} + revoked := time.Unix(1, 0) + for _, test := range []struct { + name string + secret string + tx fakeTx + want error + storageCalls []string + txCalls []string + }{ + {"claims the Environment's key", secret, fakeTx{loadSessionCreator: theCreator, lockProject: returns(false), listExecutorCredentials: returns([]ExecutorCredential{}), issueExecutorCredential: claimed}, + nil, claiming, append(checked, "IssueExecutorCredential creator "+testEnvironment+" "+testEnvironment)}, + {"retries with the same secret", secret, fakeTx{loadSessionCreator: theCreator, lockProject: returns(false), listExecutorCredentials: returns(own), authenticateExecutor: returns(true)}, + nil, claiming, append(checked, "AuthenticateExecutor "+testEnvironment+" "+digest)}, + {"another secret", secret, fakeTx{loadSessionCreator: theCreator, lockProject: returns(false), listExecutorCredentials: returns(own), authenticateExecutor: returns(false)}, + ErrExecutorCredentialExists, claiming, append(checked, "AuthenticateExecutor "+testEnvironment+" "+digest)}, + {"a revoked key", secret, fakeTx{loadSessionCreator: theCreator, lockProject: returns(false), listExecutorCredentials: returns([]ExecutorCredential{{KeyID: testEnvironment, RevokedAt: &revoked}})}, + ErrExecutorCredentialExists, claiming, checked}, + {"an operator-issued key", secret, fakeTx{loadSessionCreator: theCreator, lockProject: returns(false), listExecutorCredentials: returns([]ExecutorCredential{{KeyID: testKey}})}, + ErrExecutorCredentialExists, claiming, checked}, + {"a second key", secret, fakeTx{loadSessionCreator: theCreator, lockProject: returns(false), listExecutorCredentials: returns(append(own, ExecutorCredential{KeyID: testKey}))}, + ErrExecutorCredentialExists, claiming, checked}, + {"archived since the authorization", secret, fakeTx{loadSessionCreator: theCreator, lockProject: returns(true)}, + projects.ErrArchived, claiming, checked[:2]}, + {"another creator's Session", secret, fakeTx{loadSessionCreator: loads(identity.Subject{Kind: "user", ID: "other"}, true)}, + ErrNotFound, claiming, checked[:1]}, + {"a padded secret", secret + "=", fakeTx{}, ErrInvalidInput, validation, nil}, + {"a short secret", secret[:40], fakeTx{}, ErrInvalidInput, validation, nil}, + } { + t.Run(test.name, func(t *testing.T) { + tx := test.tx + tx.t = t + storage := &fakeStorage{t: t, installationKey: "key", getEnvironment: returns(selfHosted), loadProjectArchived: returns(false), credentials: &tx} + token := installationToken(t, storage, InstallationAuthorization{Principal: testPrincipal, Environment: testEnvironment, Version: "1.2.3", ExpiresAt: time.Now().Add(time.Minute).Unix()}) + err := deviceService(t, storage).ClaimEnvironmentInstallation(t.Context(), token, "1.2.3", test.secret) + if test.want == nil && err != nil || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v, want %v", err, test.want) + } + assertStorageCalls(t, storage, test.storageCalls...) + assertCalls(t, &tx, test.txCalls...) + }) + } + storage := &fakeStorage{t: t, installationKey: "key"} + if err := deviceService(t, storage).ClaimEnvironmentInstallation(t.Context(), "payload.signature", "1.2.3", secret); !errors.Is(err, ErrInstallationAuthorization) { + t.Fatalf("unsigned token: %v", err) + } + assertStorageCalls(t, storage, "VerifyInstallation") +} + +// A storage failure to verify, such as a missing credential key, is not an +// invalid authorization. +func TestInstallationVerificationErrorsPassThrough(t *testing.T) { + unavailable := errors.New("credential key unavailable") + for name, use := range map[string]func(*Service) error{ + "validate": func(s *Service) error { + _, err := s.ValidateEnvironmentInstallation(t.Context(), "payload.signature", "1.2.3") + return err + }, + "claim": func(s *Service) error { + return s.ClaimEnvironmentInstallation(t.Context(), "payload.signature", "1.2.3", "secret") + }, + } { + storage := &fakeStorage{t: t, verifyInstallation: unavailable} + if err := use(deviceService(t, storage)); err != unavailable { + t.Fatalf("%s: %v, want the storage error unchanged", name, err) + } + assertStorageCalls(t, storage, "VerifyInstallation") + } +} diff --git a/services/core/internal/sessions/execution_environment.go b/services/core/internal/sessions/execution_environment.go index 764fa045a..8974f4955 100644 --- a/services/core/internal/sessions/execution_environment.go +++ b/services/core/internal/sessions/execution_environment.go @@ -32,6 +32,16 @@ type EnvironmentExecution interface { // 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 + // WithFileWriteReservation reads the tenant's Environment, then runs apply + // in the transaction of its Session, with the Environment and what the + // Session lock shows. An Environment of a publicly deleted Session, or a + // missing one, is ErrNotFound. + WithFileWriteReservation(ctx context.Context, tenant, environment string, apply func(context.Context, FileWriteReservationTx, Environment, LockedSession) error) error + // WithFileWriteSettlement reads the file write to the tenant's + // Environment, including one whose Session was publicly deleted, then runs + // apply in the transaction of its Session. A malformed ID is + // ErrInvalidInput and an unknown write ErrNotFound. + WithFileWriteSettlement(ctx context.Context, tenant, environment, write string, apply func(context.Context, FileWriteSettlementTx) 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 diff --git a/services/core/internal/sessions/execution_file_writes.go b/services/core/internal/sessions/execution_file_writes.go new file mode 100644 index 000000000..3546d78af --- /dev/null +++ b/services/core/internal/sessions/execution_file_writes.go @@ -0,0 +1,171 @@ +package sessions + +import ( + "context" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/writeaudit" +) + +// FileWriteTx reads the file writes to one Environment inside its Session's +// transaction. +type FileWriteTx interface { + // LoadFileWrite reads the write with the ID to the Environment and reports + // whether it exists. + LoadFileWrite(ctx context.Context, write string) (EnvironmentFileWrite, bool, error) +} + +// FileWriteReservationTx is the Session transaction a file write is reserved +// in. +type FileWriteReservationTx interface { + FileWriteTx + ComputeAdmissionTx + InputStartTx + // 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) + // LoadPendingInput reports whether Environment input reserved for the + // Session is pending. + LoadPendingInput(ctx context.Context) (bool, error) + // CreateFileWrite stores the pending write to the Environment with the + // write-audit source in ctx, if any, and returns it. + CreateFileWrite(ctx context.Context, key FileWriteIdentity) (EnvironmentFileWrite, error) +} + +// FileWriteSettlementTx is the Session transaction a file write is settled in. +type FileWriteSettlementTx interface { + FileWriteTx + // SettleFileWrite settles the pending write in the state and returns it. + SettleFileWrite(ctx context.Context, write, state string) (EnvironmentFileWrite, error) + // RecordFileWriteAudit records the upload_file write audit of the write + // under the write-audit source stored when it was reserved. A write + // reserved without one records nothing. + RecordFileWriteAudit(ctx context.Context, write string) error +} + +// ReserveEnvironmentFileWrite persists a file write's intent before it is +// dispatched to the Environment's Runtime device. A retry with the same +// identity replays the earlier write, Replayed, and never authorizes sending +// an unknown write again; another identity for the ID is +// ErrIdempotencyConflict. The write-audit source in ctx, when present, must be +// valid for the tenant; it is stored with the intent and audited when the +// write commits. +func (o *ExecutionOperations) ReserveEnvironmentFileWrite(ctx context.Context, tenant, environment string, key FileWriteIdentity) (EnvironmentFileWrite, error) { + if !key.Valid() { + return EnvironmentFileWrite{}, ErrInvalidInput + } + var result EnvironmentFileWrite + err := o.storage.WithFileWriteReservation(ctx, tenant, environment, func(ctx context.Context, tx FileWriteReservationTx, owner Environment, locked LockedSession) error { + if source, ok := writeaudit.FromContext(ctx); ok { + if err := source.Validate(tenant); err != nil { + return err + } + } + if err := locked.Public(); err != nil { + return err + } + previous, found, err := tx.LoadFileWrite(ctx, key.ID) + if err != nil { + return err + } + if found { + if previous.Identity != key { + return ErrIdempotencyConflict + } + previous.Replayed = true + result = previous + return nil + } + if err := admitFileWrite(ctx, tx, owner, key); err != nil { + return err + } + result, err = tx.CreateFileWrite(ctx, key) + return err + }) + if err != nil { + return EnvironmentFileWrite{}, err + } + return result, nil +} + +// admitFileWrite admits a new write to the Session's live Environment through +// the device the key names, dedicated to that Environment, while the Session +// is idle with no pending input. A failed or expired Environment is +// ErrInvalidInput, a Session without a device ErrNotFound, another device +// ErrDeviceBindingConflict, and work in progress ErrTurnConflict. +func admitFileWrite(ctx context.Context, tx FileWriteReservationTx, owner Environment, key FileWriteIdentity) error { + current, err := tx.LoadEnvironment(ctx) + if err != nil { + return err + } + if current.ID != owner.ID || terminalEnvironment(current.Status) { + return ErrInvalidInput + } + bound, found, err := tx.LoadSessionDevice(ctx) + if err != nil { + return err + } + if !found { + return ErrNotFound + } + if bound.ID != key.DeviceID || bound.EnvironmentID != owner.ID { + return ErrDeviceBindingConflict + } + if err := CheckComputeAdmission(ctx, tx); err != nil { + return err + } + if err := CheckInputStart(ctx, tx); err != nil { + return err + } + pending, err := tx.LoadPendingInput(ctx) + if err != nil { + return err + } + if pending { + return ErrTurnConflict + } + return nil +} + +// SettleEnvironmentFileWrite settles a pending write as committed or rejected +// from an independently validated exact receipt. A missing receipt, +// cancellation or owner retirement is not a rejected upload, so the caller +// leaves the write pending. Settling it again in the same state replays it, +// Replayed; another identity or state is ErrIdempotencyConflict. It also +// settles a write whose Session was publicly deleted. A committed write +// records its upload_file write audit. +func (o *ExecutionOperations) SettleEnvironmentFileWrite(ctx context.Context, tenant, environment string, key FileWriteIdentity, state string) (EnvironmentFileWrite, error) { + if !key.Valid() || (state != "committed" && state != "rejected") { + return EnvironmentFileWrite{}, ErrInvalidInput + } + var result EnvironmentFileWrite + err := o.storage.WithFileWriteSettlement(ctx, tenant, environment, key.ID, func(ctx context.Context, tx FileWriteSettlementTx) error { + current, found, err := tx.LoadFileWrite(ctx, key.ID) + if err != nil { + return err + } + if !found { + return ErrNotFound + } + if current.Identity != key || (current.State != "pending" && current.State != state) { + return ErrIdempotencyConflict + } + if current.State == state { + current.Replayed = true + result = current + return nil + } + if result, err = tx.SettleFileWrite(ctx, key.ID, state); err != nil { + return err + } + if state == "committed" { + return tx.RecordFileWriteAudit(ctx, key.ID) + } + return nil + }) + if err != nil { + return EnvironmentFileWrite{}, err + } + return result, nil +} diff --git a/services/core/internal/sessions/execution_file_writes_test.go b/services/core/internal/sessions/execution_file_writes_test.go new file mode 100644 index 000000000..9fab5574f --- /dev/null +++ b/services/core/internal/sessions/execution_file_writes_test.go @@ -0,0 +1,155 @@ +package sessions + +import ( + "context" + "errors" + "testing" + "time" +) + +func (f *fakeTx) LoadFileWrite(_ context.Context, write string) (EnvironmentFileWrite, bool, error) { + f.record("LoadFileWrite", f.loadFileWrite != nil, write) + return f.loadFileWrite() +} + +func (f *fakeTx) LoadPendingInput(context.Context) (bool, error) { + f.record("LoadPendingInput", f.loadPendingInput != nil) + return f.loadPendingInput() +} + +func (f *fakeTx) CreateFileWrite(_ context.Context, key FileWriteIdentity) (EnvironmentFileWrite, error) { + f.record("CreateFileWrite", f.createFileWrite != nil, key.ID, key.DeviceID) + return f.createFileWrite() +} + +func (f *fakeTx) SettleFileWrite(_ context.Context, write, state string) (EnvironmentFileWrite, error) { + f.record("SettleFileWrite", f.settleFileWrite != nil, write, state) + return f.settleFileWrite() +} + +func (f *fakeTx) RecordFileWriteAudit(_ context.Context, write string) error { + f.record("RecordFileWriteAudit", f.recordFileWriteAudit != nil, write) + return f.recordFileWriteAudit() +} + +func (s *fakeExecutionStorage) WithFileWriteReservation(ctx context.Context, tenant, environment string, apply func(context.Context, FileWriteReservationTx, Environment, LockedSession) error) error { + s.environmentTx("WithFileWriteReservation").record("WithFileWriteReservation", true, tenant, environment) + return apply(ctx, s.tx, s.environment, s.locked) +} + +func (s *fakeExecutionStorage) WithFileWriteSettlement(ctx context.Context, tenant, environment, write string, apply func(context.Context, FileWriteSettlementTx) error) error { + s.environmentTx("WithFileWriteSettlement").record("WithFileWriteSettlement", true, tenant, environment, write) + return apply(ctx, s.tx) +} + +const ( + testFileWrite = "7a1e4c2b-3d5f-4a6b-8c9d-0e1f2a3b4c5d" + testDevice = "5b0cf2a4-6c41-4f55-9d2f-3f6f1b1c7e10" +) + +var testWriteKey = FileWriteIdentity{ID: testFileWrite, DeviceID: testDevice, RequestSHA256: credentialDigest} + +func TestReserveEnvironmentFileWrite(t *testing.T) { + reserve := "WithFileWriteReservation " + testTenant + " " + testEnvironment + admission := []string{reserve, "LoadFileWrite " + testFileWrite, "LoadEnvironment", "LoadSessionDevice", "LoadComputeSuspension", "LoadPendingFileWrite", "LoadActiveTurn", "LoadPendingInput"} + pending := EnvironmentFileWrite{Identity: testWriteKey, EnvironmentID: testEnvironment, SessionID: testSession, State: "pending"} + bound := loads(ExecutionDevice{ID: testDevice, EnvironmentID: testEnvironment}, true) + admitted := func(tx fakeTx) fakeTx { + tx.loadFileWrite = loads(EnvironmentFileWrite{}, false) + if tx.loadEnvironment == nil { + tx.loadEnvironment = returns(selfHosted) + } + if tx.loadSessionDevice == nil { + tx.loadSessionDevice = bound + } + return tx + } + idle := func(tx fakeTx) fakeTx { + tx = admitted(tx) + tx.loadComputeSuspension, tx.loadPendingFileWrite, tx.loadActiveTurn = returns(false), returns(false), activeTurn(nil) + if tx.loadPendingInput == nil { + tx.loadPendingInput = returns(false) + } + return tx + } + other := EnvironmentFileWrite{Identity: FileWriteIdentity{ID: testFileWrite, DeviceID: testKey, RequestSHA256: credentialDigest}, State: "pending"} + for _, test := range []struct { + name string + tx fakeTx + locked LockedSession + want error + replayed bool + calls []string + }{ + {"reserves", idle(fakeTx{createFileWrite: returns(pending)}), LockedSession{}, nil, false, append(admission, "CreateFileWrite "+testFileWrite+" "+testDevice)}, + {"replays the same identity", fakeTx{loadFileWrite: loads(pending, true)}, LockedSession{}, nil, true, admission[:2]}, + {"another identity", fakeTx{loadFileWrite: loads(other, true)}, LockedSession{}, ErrIdempotencyConflict, false, admission[:2]}, + {"a deleted Session", fakeTx{}, LockedSession{Deleted: true}, ErrNotFound, false, admission[:1]}, + {"a replaced Environment", admitted(fakeTx{loadEnvironment: returns(Environment{ID: testKey, Status: "ready"})}), LockedSession{}, ErrInvalidInput, false, admission[:3]}, + {"a failed Environment", admitted(fakeTx{loadEnvironment: returns(Environment{ID: testEnvironment, Status: "failed"})}), LockedSession{}, ErrInvalidInput, false, admission[:3]}, + {"no device", admitted(fakeTx{loadSessionDevice: loads(ExecutionDevice{}, false)}), LockedSession{}, ErrNotFound, false, admission[:4]}, + {"another device", admitted(fakeTx{loadSessionDevice: loads(ExecutionDevice{ID: testKey, EnvironmentID: testEnvironment}, true)}), LockedSession{}, ErrDeviceBindingConflict, false, admission[:4]}, + {"suspended compute", func() fakeTx { tx := admitted(fakeTx{}); tx.loadComputeSuspension = returns(true); return tx }(), LockedSession{}, ErrTurnConflict, false, admission[:5]}, + {"pending input", idle(fakeTx{loadPendingInput: returns(true)}), LockedSession{}, ErrTurnConflict, false, admission}, + } { + t.Run(test.name, func(t *testing.T) { + tx := test.tx + tx.t = t + storage := &fakeExecutionStorage{tx: &tx, environment: selfHosted, locked: test.locked} + write, err := executionOperations(t, storage).ReserveEnvironmentFileWrite(t.Context(), testTenant, testEnvironment, testWriteKey) + if test.want == nil && err != nil || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v, want %v", err, test.want) + } + if test.want == nil && (write.Identity != testWriteKey || write.Replayed != test.replayed) { + t.Fatalf("write %+v", write) + } + assertCalls(t, &tx, test.calls...) + }) + } + if _, err := executionOperations(t, &fakeExecutionStorage{}).ReserveEnvironmentFileWrite(t.Context(), testTenant, testEnvironment, FileWriteIdentity{ID: testFileWrite, DeviceID: testDevice}); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("no request digest: %v", err) + } +} + +func TestSettleEnvironmentFileWrite(t *testing.T) { + settle := "WithFileWriteSettlement " + testTenant + " " + testEnvironment + " " + testFileWrite + load := "LoadFileWrite " + testFileWrite + settledAt := time.Unix(1, 0) + write := func(state string) EnvironmentFileWrite { + return EnvironmentFileWrite{Identity: testWriteKey, EnvironmentID: testEnvironment, SessionID: testSession, State: state, SettledAt: &settledAt} + } + for _, test := range []struct { + name string + state string + tx fakeTx + want error + replayed bool + calls []string + }{ + {"commits and audits", "committed", fakeTx{loadFileWrite: loads(write("pending"), true), settleFileWrite: returns(write("committed")), recordFileWriteAudit: done}, + nil, false, []string{settle, load, "SettleFileWrite " + testFileWrite + " committed", "RecordFileWriteAudit " + testFileWrite}}, + {"rejects without an audit", "rejected", fakeTx{loadFileWrite: loads(write("pending"), true), settleFileWrite: returns(write("rejected"))}, + nil, false, []string{settle, load, "SettleFileWrite " + testFileWrite + " rejected"}}, + {"replays the same state", "committed", fakeTx{loadFileWrite: loads(write("committed"), true)}, nil, true, []string{settle, load}}, + {"another state", "rejected", fakeTx{loadFileWrite: loads(write("committed"), true)}, ErrIdempotencyConflict, false, []string{settle, load}}, + {"another identity", "committed", fakeTx{loadFileWrite: loads(EnvironmentFileWrite{Identity: FileWriteIdentity{ID: testFileWrite, DeviceID: testKey, RequestSHA256: credentialDigest}, State: "pending"}, true)}, + ErrIdempotencyConflict, false, []string{settle, load}}, + {"an unknown write", "committed", fakeTx{loadFileWrite: loads(EnvironmentFileWrite{}, false)}, ErrNotFound, false, []string{settle, load}}, + } { + t.Run(test.name, func(t *testing.T) { + tx := test.tx + tx.t = t + settled, err := executionOperations(t, &fakeExecutionStorage{tx: &tx}).SettleEnvironmentFileWrite(t.Context(), testTenant, testEnvironment, testWriteKey, test.state) + if test.want == nil && err != nil || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v, want %v", err, test.want) + } + if test.want == nil && (settled.State != test.state || settled.Replayed != test.replayed) { + t.Fatalf("settled %+v", settled) + } + assertCalls(t, &tx, test.calls...) + }) + } + if _, err := executionOperations(t, &fakeExecutionStorage{}).SettleEnvironmentFileWrite(t.Context(), testTenant, testEnvironment, testWriteKey, "pending"); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("pending state: %v", err) + } +} diff --git a/services/core/internal/sessions/execution_functions_test.go b/services/core/internal/sessions/execution_functions_test.go index f3c016138..42531197d 100644 --- a/services/core/internal/sessions/execution_functions_test.go +++ b/services/core/internal/sessions/execution_functions_test.go @@ -21,6 +21,8 @@ type fakeExecutionStorage struct { tx *fakeTx locked LockedSession + // environment is the Environment WithFileWriteReservation reads. + environment Environment // pages are the pages ListEnvironmentConnections reads, one per call. pages [][]EnvironmentKey // connections is the error each WithConnection call returns in place of diff --git a/services/core/internal/sessions/executor_credentials.go b/services/core/internal/sessions/executor_credentials.go index c88505295..dcde3026d 100644 --- a/services/core/internal/sessions/executor_credentials.go +++ b/services/core/internal/sessions/executor_credentials.go @@ -2,18 +2,32 @@ package sessions import ( "context" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "fmt" "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" ) -// ExecutorCredentialReader authenticates Environment executor credentials. +// ExecutorCredentialReader authenticates Environment executor credentials and +// reads a Project's credentials for an Environment. 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) + // ProjectExecutorCredentialState reads, in one snapshot, the connection + // facts of the tenant's self_hosted Environment and the metadata of the + // credentials restricted to it that the Project's principal holds, oldest + // first. An invalid principal is ErrInvalidInput; an Environment of a + // publicly deleted Session, another type or another tenant is ErrNotFound. + // The snapshot ends before any live observation of the connection. + ProjectExecutorCredentialState(ctx context.Context, project identity.Principal, environment string) (ExecutorCredentialState, error) } type IssuedExecutorCredential struct { @@ -47,12 +61,303 @@ type ExecutorCredentialState struct { Connection ExecutorConnectionState } -// InstallationAuthorization permits claiming one Environment's connect-only key. -// The Environment UUID is reserved as that key's ID. Reissuing an authorization -// never rotates or revives the key, and a retry must prove the same local secret. -type InstallationAuthorization struct { - Principal identity.Principal `json:"principal"` - Environment string `json:"environment_id"` - Version string `json:"version"` - ExpiresAt int64 `json:"expires_at"` +// ExecutorCredentialGrant is a new executor credential: its key, the principal +// it executes as, the Environment it is restricted to, empty for none, and the +// SHA-256 digest of its secret. +type ExecutorCredentialGrant struct { + Principal identity.Principal + KeyID string + EnvironmentID string + Digest string +} + +// ExecutorCredentialStorage stores Environment executor credentials and signs +// installation authorizations. +type ExecutorCredentialStorage interface { + // GetEnvironment reads the tenant's Environment of a Session that was not + // publicly deleted, as EnvironmentReader does. + GetEnvironment(ctx context.Context, tenant, environment string) (Environment, error) + // LoadProjectArchived reports whether the tenant's Project is archived. A + // tenant without a Project is ErrNotFound. + LoadProjectArchived(ctx context.Context, tenant string) (bool, error) + // ExecutorProjectExists reports whether the scope's organization and + // Project map to its execution tenant. + ExecutorProjectExists(ctx context.Context, scope identity.ProjectScope) (bool, error) + // LoadExecutorCredentialRestriction reads the Environment the principal's + // key is restricted to, empty for an unrestricted key. A key of another + // principal or Project, or an unknown one, is ErrNotFound. + LoadExecutorCredentialRestriction(ctx context.Context, principal identity.Principal, key string) (string, error) + // SignInstallation returns the signature of an installation authorization + // payload under the credential key. Without that key it is + // credentialcrypto.ErrUnavailable. + SignInstallation(ctx context.Context, payload string) (string, error) + // VerifyInstallation checks the signature of an installation + // authorization payload: another signature is + // ErrInstallationAuthorization, and a service without the credential key + // credentialcrypto.ErrUnavailable. + VerifyInstallation(ctx context.Context, payload, signature string) error + // WithExecutorCredentials runs apply in a transaction of the tenant. + WithExecutorCredentials(ctx context.Context, tenant string, apply func(context.Context, ExecutorCredentialTx) error) error + // WithEnvironmentExecutorCredentials reads the tenant's Environment, then + // runs apply in the transaction of its Session, with what the Session lock + // shows. A malformed Environment ID is ErrInvalidInput, and an Environment + // of a publicly deleted Session, or a missing one, ErrNotFound. + WithEnvironmentExecutorCredentials(ctx context.Context, tenant, environment string, apply func(context.Context, EnvironmentExecutorCredentialTx, LockedSession) error) error +} + +// ExecutorCredentialTx is the transaction of one tenant that executor +// credentials are written in. +type ExecutorCredentialTx interface { + // LockProject share-locks the tenant's Project, which archiving updates, + // and reports whether it is archived, so a write either commits before the + // archive or sees it. A tenant without a Project is ErrNotFound. + LockProject(ctx context.Context) (bool, error) + // ListExecutorCredentials lists the subject's keys restricted to the + // Environment, oldest first. + ListExecutorCredentials(ctx context.Context, environment string, subject identity.Subject) ([]ExecutorCredential, error) + // AuthenticateExecutor reports whether the secret with the SHA-256 digest + // authenticates for the Environment, as AuthenticateEnvironmentExecutor + // checks. + AuthenticateExecutor(ctx context.Context, environment, digest string) (bool, error) + // IssueExecutorCredential stores the new credential in its principal's + // Project and returns its key and restriction. When no credential is + // stored, because the key ID exists or the restriction names an + // Environment the principal did not create, it is + // ErrExecutorCredentialExists. + IssueExecutorCredential(ctx context.Context, grant ExecutorCredentialGrant) (IssuedExecutorCredential, error) + // RotateExecutorCredential replaces the secret of the subject's key with + // the one whose SHA-256 digest is given, restores the key if it was + // revoked, and returns its key and restriction. An unknown key is + // ErrNotFound. + RotateExecutorCredential(ctx context.Context, subject identity.Subject, key, digest string) (IssuedExecutorCredential, error) + // RevokeExecutorCredential revokes the subject's key; revoking it again + // changes nothing. An unknown key is ErrNotFound. + RevokeExecutorCredential(ctx context.Context, subject identity.Subject, key string) error + // RecordExecutorCredentialAudit records the administrator audit entry of + // the action on the key under the provenance in ctx. It never records a + // secret. + RecordExecutorCredentialAudit(ctx context.Context, action, key string) error +} + +// EnvironmentExecutorCredentialTx is the transaction of the Session whose +// Environment a credential is restricted to. +type EnvironmentExecutorCredentialTx interface { + ExecutorCredentialTx + // LoadSessionCreator reads who created the Session and reports whether + // the Session recorded a creator. + LoadSessionCreator(ctx context.Context) (identity.Subject, bool, error) +} + +// executorSecret is a new connect-only secret and the SHA-256 digest that is +// stored instead of it. +type executorSecret struct{ token, digest string } + +// newExecutorSecret generates 32 random bytes as an unpadded base64url token. +func newExecutorSecret() (executorSecret, error) { + secret := make([]byte, 32) + if _, err := rand.Read(secret); err != nil { + return executorSecret{}, err + } + token := base64.RawURLEncoding.EncodeToString(secret) + return executorSecret{token: token, digest: executorDigest(token)}, nil +} + +func executorDigest(token string) string { + hash := sha256.Sum256([]byte(token)) + return hex.EncodeToString(hash[:]) +} + +// installationSecretDigest returns the digest of a secret an installer +// generated: exactly 32 bytes in canonical unpadded base64url. Anything else is +// ErrInvalidInput. +func installationSecretDigest(secret string) (string, error) { + decoded, err := base64.RawURLEncoding.DecodeString(secret) + if err != nil || len(decoded) != 32 || base64.RawURLEncoding.EncodeToString(decoded) != secret { + return "", ErrInvalidInput + } + return executorDigest(secret), nil +} + +// checkExecutorIdentity validates the principal a credential executes as and +// the credential's key ID. +func checkExecutorIdentity(principal identity.Principal, key string) error { + if err := principal.Validate(); err != nil { + return fmt.Errorf("%w: %v", ErrInvalidInput, err) + } + if !validID(key) { + return ErrInvalidInput + } + return nil +} + +// activeProject is projects.ErrArchived for an archived Project. +func activeProject(archived bool) error { + if archived { + return projects.ErrArchived + } + return nil +} + +// lockActiveProject repeats the archive check in the writing transaction. +func lockActiveProject(ctx context.Context, tx ExecutorCredentialTx) error { + archived, err := tx.LockProject(ctx) + if err != nil { + return err + } + return activeProject(archived) +} + +// checkActiveProject reads, before any transaction, whether the tenant's +// Project is archived. +func (s *Service) checkActiveProject(ctx context.Context, tenant string) error { + archived, err := s.storage.LoadProjectArchived(ctx, tenant) + if err != nil { + return err + } + return activeProject(archived) +} + +// selfHostedTarget reads the principal's self_hosted Environment an executor +// credential or installation targets. Another type is ErrNotFound. +func (s *Service) selfHostedTarget(ctx context.Context, principal identity.Principal, environment string) (Environment, error) { + if err := principal.Validate(); err != nil { + return Environment{}, ErrInvalidInput + } + current, err := s.storage.GetEnvironment(ctx, principal.TenantID, environment) + if err != nil { + return Environment{}, err + } + if kind, err := EnvironmentType(current.Configuration); err != nil || kind != "self_hosted" { + return Environment{}, ErrNotFound + } + return current, nil +} + +// withExecutorTarget runs apply in the transaction that writes a credential +// restricted to environment, or unrestricted when it is empty. A restricted +// credential is written in the transaction of the Environment's Session, whose +// lock orders it against the Session's deletion, and only by the Session's +// creator; any other Session is ErrNotFound. +func (s *Service) withExecutorTarget(ctx context.Context, principal identity.Principal, environment string, apply func(context.Context, ExecutorCredentialTx) error) error { + if environment == "" { + return s.storage.WithExecutorCredentials(ctx, principal.TenantID, apply) + } + return s.storage.WithEnvironmentExecutorCredentials(ctx, principal.TenantID, environment, func(ctx context.Context, tx EnvironmentExecutorCredentialTx, locked LockedSession) error { + if err := locked.Public(); err != nil { + return err + } + creator, recorded, err := tx.LoadSessionCreator(ctx) + if err != nil { + return err + } + if !recorded || creator != principal.Subject() { + return ErrNotFound + } + return apply(ctx, tx) + }) +} + +// IssueExecutorCredential returns a new connect-only secret once, restricted +// to the Environment unless it is empty, without replacing an existing key: +// that is ErrExecutorCredentialExists. The principal's Project must be mapped +// to its tenant. +func (s *Service) IssueExecutorCredential(ctx context.Context, principal identity.Principal, key, environment string) (IssuedExecutorCredential, error) { + return s.issueExecutorCredential(ctx, principal, key, environment, nil) +} + +// issueExecutorCredential runs before, when given, first in the issuing +// transaction; its error aborts the issuance. +func (s *Service) issueExecutorCredential(ctx context.Context, principal identity.Principal, key, environment string, before func(context.Context, ExecutorCredentialTx) error) (IssuedExecutorCredential, error) { + if err := checkExecutorIdentity(principal, key); err != nil { + return IssuedExecutorCredential{}, err + } + exists, err := s.storage.ExecutorProjectExists(ctx, principal.ProjectScope) + if err != nil { + return IssuedExecutorCredential{}, err + } + if !exists { + return IssuedExecutorCredential{}, ErrNotFound + } + if environment != "" && !validID(environment) { + return IssuedExecutorCredential{}, ErrInvalidInput + } + secret, err := newExecutorSecret() + if err != nil { + return IssuedExecutorCredential{}, err + } + var result IssuedExecutorCredential + err = s.withExecutorTarget(ctx, principal, environment, func(ctx context.Context, tx ExecutorCredentialTx) error { + if before != nil { + if err := before(ctx, tx); err != nil { + return err + } + } + issued, err := tx.IssueExecutorCredential(ctx, ExecutorCredentialGrant{Principal: principal, KeyID: key, EnvironmentID: environment, Digest: secret.digest}) + if err != nil { + return err + } + issued.Token = secret.token + result = issued + return nil + }) + if err != nil { + return IssuedExecutorCredential{}, err + } + return result, nil +} + +// RotateExecutorCredential returns a new secret for the principal's key once, +// keeping its restriction and restoring it if it was revoked. An unknown key +// is ErrNotFound. +func (s *Service) RotateExecutorCredential(ctx context.Context, principal identity.Principal, key string) (IssuedExecutorCredential, error) { + return s.rotateExecutorCredential(ctx, principal, key, nil) +} + +// rotateExecutorCredential runs before, when given, first in the rotating +// transaction; its error aborts the rotation. +func (s *Service) rotateExecutorCredential(ctx context.Context, principal identity.Principal, key string, before func(context.Context, ExecutorCredentialTx) error) (IssuedExecutorCredential, error) { + if err := checkExecutorIdentity(principal, key); err != nil { + return IssuedExecutorCredential{}, err + } + restriction, err := s.storage.LoadExecutorCredentialRestriction(ctx, principal, key) + if err != nil { + return IssuedExecutorCredential{}, err + } + secret, err := newExecutorSecret() + if err != nil { + return IssuedExecutorCredential{}, err + } + var result IssuedExecutorCredential + err = s.withExecutorTarget(ctx, principal, restriction, func(ctx context.Context, tx ExecutorCredentialTx) error { + if before != nil { + if err := before(ctx, tx); err != nil { + return err + } + } + rotated, err := tx.RotateExecutorCredential(ctx, principal.Subject(), key, secret.digest) + if err != nil { + return err + } + rotated.Token = secret.token + result = rotated + return nil + }) + if err != nil { + return IssuedExecutorCredential{}, err + } + return result, nil +} + +// RevokeExecutorCredential revokes the principal's key; revoking it again +// succeeds. An unknown key is ErrNotFound. +func (s *Service) RevokeExecutorCredential(ctx context.Context, principal identity.Principal, key string) error { + if err := checkExecutorIdentity(principal, key); err != nil { + return err + } + if _, err := s.storage.LoadExecutorCredentialRestriction(ctx, principal, key); err != nil { + return err + } + return s.storage.WithExecutorCredentials(ctx, principal.TenantID, func(ctx context.Context, tx ExecutorCredentialTx) error { + return tx.RevokeExecutorCredential(ctx, principal.Subject(), key) + }) } diff --git a/services/core/internal/sessions/executor_credentials_test.go b/services/core/internal/sessions/executor_credentials_test.go new file mode 100644 index 000000000..fb8964205 --- /dev/null +++ b/services/core/internal/sessions/executor_credentials_test.go @@ -0,0 +1,288 @@ +package sessions + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "strings" + "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" +) + +func (s *fakeStorage) GetEnvironment(_ context.Context, tenant, environment string) (Environment, error) { + s.record("GetEnvironment", s.getEnvironment != nil, tenant, environment) + return s.getEnvironment() +} + +func (s *fakeStorage) LoadProjectArchived(_ context.Context, tenant string) (bool, error) { + s.record("LoadProjectArchived", s.loadProjectArchived != nil, tenant) + return s.loadProjectArchived() +} + +func (s *fakeStorage) ExecutorProjectExists(_ context.Context, scope identity.ProjectScope) (bool, error) { + s.record("ExecutorProjectExists", s.projectExists != nil, scope.TenantID) + return s.projectExists() +} + +func (s *fakeStorage) LoadExecutorCredentialRestriction(_ context.Context, principal identity.Principal, key string) (string, error) { + s.record("LoadExecutorCredentialRestriction", s.loadRestriction != nil, principal.TenantID, key) + return s.loadRestriction() +} + +// signature is the fake signature of an installation payload. +func (s *fakeStorage) signature(payload string) string { + sum := sha256.Sum256([]byte(s.installationKey + "." + payload)) + return hex.EncodeToString(sum[:]) +} + +func (s *fakeStorage) SignInstallation(_ context.Context, payload string) (string, error) { + s.record("SignInstallation", s.installationKey != "") + return s.signature(payload), nil +} + +func (s *fakeStorage) VerifyInstallation(_ context.Context, payload, signature string) error { + s.record("VerifyInstallation", s.installationKey != "" || s.verifyInstallation != nil) + if s.verifyInstallation != nil { + return s.verifyInstallation + } + if signature != s.signature(payload) { + return ErrInstallationAuthorization + } + return nil +} + +func (s *fakeStorage) WithExecutorCredentials(ctx context.Context, tenant string, apply func(context.Context, ExecutorCredentialTx) error) error { + s.record("WithExecutorCredentials", s.credentials != nil, tenant) + return apply(ctx, s.credentials) +} + +func (s *fakeStorage) WithEnvironmentExecutorCredentials(ctx context.Context, tenant, environment string, apply func(context.Context, EnvironmentExecutorCredentialTx, LockedSession) error) error { + s.record("WithEnvironmentExecutorCredentials", s.credentials != nil, tenant, environment) + return apply(ctx, s.credentials, s.locked) +} + +func (f *fakeTx) LockProject(context.Context) (bool, error) { + f.record("LockProject", f.lockProject != nil) + return f.lockProject() +} + +func (f *fakeTx) ListExecutorCredentials(_ context.Context, environment string, subject identity.Subject) ([]ExecutorCredential, error) { + f.record("ListExecutorCredentials", f.listExecutorCredentials != nil, environment, subject.ID) + return f.listExecutorCredentials() +} + +func (f *fakeTx) AuthenticateExecutor(_ context.Context, environment, digest string) (bool, error) { + f.record("AuthenticateExecutor", f.authenticateExecutor != nil, environment, digest) + return f.authenticateExecutor() +} + +func (f *fakeTx) IssueExecutorCredential(_ context.Context, grant ExecutorCredentialGrant) (IssuedExecutorCredential, error) { + f.record("IssueExecutorCredential", f.issueExecutorCredential != nil, grant.Principal.SubjectID, grant.KeyID, grant.EnvironmentID) + return f.issueExecutorCredential(grant) +} + +func (f *fakeTx) RotateExecutorCredential(_ context.Context, subject identity.Subject, key, digest string) (IssuedExecutorCredential, error) { + f.record("RotateExecutorCredential", f.rotateExecutorCredential != nil, subject.ID, key) + return f.rotateExecutorCredential(digest) +} + +func (f *fakeTx) RevokeExecutorCredential(_ context.Context, subject identity.Subject, key string) error { + f.record("RevokeExecutorCredential", f.revokeExecutorCredential != nil, subject.ID, key) + return f.revokeExecutorCredential() +} + +func (f *fakeTx) RecordExecutorCredentialAudit(_ context.Context, action, key string) error { + f.record("RecordExecutorCredentialAudit", f.recordExecutorCredentialAudit != nil, action, key) + return f.recordExecutorCredentialAudit() +} + +func (f *fakeTx) LoadSessionCreator(context.Context) (identity.Subject, bool, error) { + f.record("LoadSessionCreator", f.loadSessionCreator != nil) + return f.loadSessionCreator() +} + +// fails is a fake method that fails with err. +func fails[T any](err error) func() (T, error) { + return func() (T, error) { + var zero T + return zero, err + } +} + +const ( + testKey = "0d9a4f3e-5b1c-4c2e-8f6a-7e3b2d1c0a99" + testEnvironment = "2b5e8c1d-4f3a-4e6b-9c7d-8a1f0e2d3c4b" +) + +var ( + testPrincipal = identity.Principal{ProjectScope: identity.ProjectScope{TenantID: testTenant, OrganizationID: "organization", ProjectID: "project"}, SubjectKind: "user", SubjectID: "creator"} + selfHosted = Environment{ID: testEnvironment, SessionID: testSession, Status: "ready", Configuration: []byte(`{"type":"self_hosted"}`)} + // theCreator is a fake LoadSessionCreator: testPrincipal created the + // Session. + theCreator = loads(testPrincipal.Subject(), true) +) + +// credentialCase is one executor credential use case: the storage and +// transaction it runs against, what it returns and the calls it makes. +type credentialCase struct { + name string + storage fakeStorage + tx fakeTx + operate func(*Service) (IssuedExecutorCredential, error) + want error + storageCalls []string + txCalls []string +} + +// runCredentialCases runs each case and checks its error, calls and issued secret: a +// successful issuance or rotation returns the token whose digest it stored. +func runCredentialCases(t *testing.T, cases []credentialCase, digest *string) { + t.Helper() + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + *digest = "" + tx := test.tx + tx.t = t + storage := test.storage + storage.t, storage.credentials = t, &tx + service, err := NewService(&storage) + if err != nil { + t.Fatal(err) + } + issued, err := test.operate(service) + if test.want == nil && err != nil || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v, want %v", err, test.want) + } + assertStorageCalls(t, &storage, test.storageCalls...) + assertCalls(t, &tx, test.txCalls...) + if (issued.Token != "" || *digest != "") && (err != nil || executorDigest(issued.Token) != *digest) { + t.Fatalf("issued %+v with digest %q, %v", issued, *digest, err) + } + }) + } +} + +func TestExecutorCredentialUseCases(t *testing.T) { + var digest string + issued := func(grant ExecutorCredentialGrant) (IssuedExecutorCredential, error) { + digest = grant.Digest + return IssuedExecutorCredential{KeyID: grant.KeyID, EnvironmentID: grant.EnvironmentID}, nil + } + rotated := func(stored string) (IssuedExecutorCredential, error) { + digest = stored + return IssuedExecutorCredential{KeyID: testKey}, nil + } + issue := func(environment string) func(*Service) (IssuedExecutorCredential, error) { + return func(s *Service) (IssuedExecutorCredential, error) { + return s.IssueExecutorCredential(t.Context(), testPrincipal, testKey, environment) + } + } + rotate := func(s *Service) (IssuedExecutorCredential, error) { + return s.RotateExecutorCredential(t.Context(), testPrincipal, testKey) + } + revoke := func(s *Service) (IssuedExecutorCredential, error) { + return IssuedExecutorCredential{}, s.RevokeExecutorCredential(t.Context(), testPrincipal, testKey) + } + exists := "ExecutorProjectExists " + testTenant + restriction := "LoadExecutorCredentialRestriction " + testTenant + " " + testKey + unrestricted := "WithExecutorCredentials " + testTenant + restricted := "WithEnvironmentExecutorCredentials " + testTenant + " " + testEnvironment + issuance := func(environment string) string { + return strings.Join([]string{"IssueExecutorCredential creator", testKey, environment}, " ") + } + rotation := "RotateExecutorCredential creator " + testKey + mapped := fakeStorage{projectExists: returns(true)} + runCredentialCases(t, []credentialCase{ + {"issues an unrestricted key", mapped, fakeTx{issueExecutorCredential: issued}, issue(""), nil, + []string{exists, unrestricted}, []string{issuance("")}}, + {"issues a key in the creator's Session", mapped, fakeTx{loadSessionCreator: theCreator, issueExecutorCredential: issued}, issue(testEnvironment), nil, + []string{exists, restricted}, []string{"LoadSessionCreator", issuance(testEnvironment)}}, + {"another creator's Session", mapped, fakeTx{loadSessionCreator: loads(identity.Subject{Kind: "user", ID: "other"}, true)}, issue(testEnvironment), ErrNotFound, + []string{exists, restricted}, []string{"LoadSessionCreator"}}, + {"a Session without a creator", mapped, fakeTx{loadSessionCreator: loads(identity.Subject{}, false)}, issue(testEnvironment), ErrNotFound, + []string{exists, restricted}, []string{"LoadSessionCreator"}}, + {"a deleted Session", fakeStorage{projectExists: returns(true), locked: LockedSession{Deleted: true}}, fakeTx{}, issue(testEnvironment), ErrNotFound, + []string{exists, restricted}, nil}, + {"an unmapped Project", fakeStorage{projectExists: returns(false)}, fakeTx{}, issue(""), ErrNotFound, []string{exists}, nil}, + {"a malformed Environment", mapped, fakeTx{}, issue("environment"), ErrInvalidInput, []string{exists}, nil}, + {"an existing key", mapped, fakeTx{issueExecutorCredential: func(ExecutorCredentialGrant) (IssuedExecutorCredential, error) { + return IssuedExecutorCredential{}, ErrExecutorCredentialExists + }}, issue(""), ErrExecutorCredentialExists, []string{exists, unrestricted}, []string{issuance("")}}, + {"rotates in the restricted Environment's Session", fakeStorage{loadRestriction: returns(testEnvironment)}, fakeTx{loadSessionCreator: theCreator, rotateExecutorCredential: rotated}, rotate, nil, + []string{restriction, restricted}, []string{"LoadSessionCreator", rotation}}, + {"rotates an unrestricted key", fakeStorage{loadRestriction: returns("")}, fakeTx{rotateExecutorCredential: rotated}, rotate, nil, + []string{restriction, unrestricted}, []string{rotation}}, + {"rotates an unknown key", fakeStorage{loadRestriction: fails[string](ErrNotFound)}, fakeTx{}, rotate, ErrNotFound, []string{restriction}, nil}, + {"revokes", fakeStorage{loadRestriction: returns(testEnvironment)}, fakeTx{revokeExecutorCredential: done}, revoke, nil, + []string{restriction, unrestricted}, []string{"RevokeExecutorCredential creator " + testKey}}, + {"revokes an unknown key", fakeStorage{loadRestriction: fails[string](ErrNotFound)}, fakeTx{}, revoke, ErrNotFound, []string{restriction}, nil}, + }, &digest) + + // A malformed principal or key is rejected before storage. + service := deviceService(t, &fakeStorage{t: t}) + invalid := testPrincipal + invalid.SubjectKind = "robot" + if _, err := service.IssueExecutorCredential(t.Context(), invalid, testKey, ""); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("invalid principal: %v", err) + } + if _, err := service.RotateExecutorCredential(t.Context(), testPrincipal, "key"); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("malformed key: %v", err) + } + if err := service.RevokeExecutorCredential(t.Context(), testPrincipal, "00000000-0000-0000-0000-000000000000"); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("nil key: %v", err) + } +} + +func TestProjectExecutorCredentialUseCases(t *testing.T) { + var digest string + issued := func(grant ExecutorCredentialGrant) (IssuedExecutorCredential, error) { + digest = grant.Digest + return IssuedExecutorCredential{KeyID: grant.KeyID, EnvironmentID: grant.EnvironmentID}, nil + } + rotated := func(stored string) (IssuedExecutorCredential, error) { + digest = stored + return IssuedExecutorCredential{KeyID: testKey, EnvironmentID: testEnvironment}, nil + } + operate := func(rotate bool) func(*Service) (IssuedExecutorCredential, error) { + return func(s *Service) (IssuedExecutorCredential, error) { + return s.IssueProjectExecutorCredential(t.Context(), testPrincipal, testEnvironment, testKey, rotate) + } + } + revoke := func(s *Service) (IssuedExecutorCredential, error) { + return IssuedExecutorCredential{}, s.RevokeProjectExecutorCredential(t.Context(), testPrincipal, testEnvironment, testKey) + } + target := "GetEnvironment " + testTenant + " " + testEnvironment + archive := "LoadProjectArchived " + testTenant + restriction := "LoadExecutorCredentialRestriction " + testTenant + " " + testKey + restricted := "WithEnvironmentExecutorCredentials " + testTenant + " " + testEnvironment + active := fakeStorage{getEnvironment: returns(selfHosted), loadProjectArchived: returns(false), projectExists: returns(true), loadRestriction: returns(testEnvironment)} + withStorage := func(change func(*fakeStorage)) fakeStorage { + storage := active + change(&storage) + return storage + } + runCredentialCases(t, []credentialCase{ + {"issues and audits", active, fakeTx{loadSessionCreator: theCreator, lockProject: returns(false), recordExecutorCredentialAudit: done, issueExecutorCredential: issued}, operate(false), nil, + []string{target, archive, "ExecutorProjectExists " + testTenant, restricted}, + []string{"LoadSessionCreator", "LockProject", "RecordExecutorCredentialAudit issue " + testKey, "IssueExecutorCredential creator " + testKey + " " + testEnvironment}}, + {"rotates and audits", active, fakeTx{loadSessionCreator: theCreator, lockProject: returns(false), recordExecutorCredentialAudit: done, rotateExecutorCredential: rotated}, operate(true), nil, + []string{target, archive, restriction, restriction, restricted}, + []string{"LoadSessionCreator", "LockProject", "RecordExecutorCredentialAudit rotate " + testKey, "RotateExecutorCredential creator " + testKey}}, + {"an OpenAI-hosted Environment", withStorage(func(s *fakeStorage) { s.getEnvironment = returns(hostedEnvironment) }), fakeTx{}, operate(false), ErrNotFound, + []string{target}, nil}, + {"an archived Project", withStorage(func(s *fakeStorage) { s.loadProjectArchived = returns(true) }), fakeTx{}, operate(false), projects.ErrArchived, + []string{target, archive}, nil}, + {"archived before the write", active, fakeTx{loadSessionCreator: theCreator, lockProject: returns(true)}, operate(false), projects.ErrArchived, + []string{target, archive, "ExecutorProjectExists " + testTenant, restricted}, []string{"LoadSessionCreator", "LockProject"}}, + {"rotates a key restricted elsewhere", withStorage(func(s *fakeStorage) { s.loadRestriction = returns("") }), fakeTx{}, operate(true), ErrNotFound, + []string{target, archive, restriction}, nil}, + {"revokes and audits in an archived Project", withStorage(func(s *fakeStorage) { s.loadProjectArchived = nil }), fakeTx{revokeExecutorCredential: done, recordExecutorCredentialAudit: done}, revoke, nil, + []string{target, restriction, "WithExecutorCredentials " + testTenant}, + []string{"RevokeExecutorCredential creator " + testKey, "RecordExecutorCredentialAudit revoke " + testKey}}, + {"revokes a key restricted elsewhere", withStorage(func(s *fakeStorage) { s.loadRestriction = returns("6c3b2a19-8d7e-4f60-a5b4-c3d2e1f0a9b8") }), fakeTx{}, revoke, ErrNotFound, + []string{target, restriction}, nil}, + }, &digest) +} diff --git a/services/core/internal/sessions/project_executor_credentials.go b/services/core/internal/sessions/project_executor_credentials.go new file mode 100644 index 000000000..60c562b2e --- /dev/null +++ b/services/core/internal/sessions/project_executor_credentials.go @@ -0,0 +1,84 @@ +package sessions + +import ( + "context" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" +) + +// The Project executor credential use cases serve Core-key credential +// management. project is the Project's own principal, which becomes the +// credential's execution principal; the target must be a self_hosted +// Environment of that Project whose Session is not deleted, and only +// credentials restricted to it are managed. Each write records an +// administrator audit entry in its transaction, under the provenance in ctx. + +// IssueProjectExecutorCredential issues a new key restricted to the +// Environment or, with rotate, replaces the secret of an existing key +// restricted to it. The target is checked first, then the archived Project +// (projects.ErrArchived), then the key. +func (s *Service) IssueProjectExecutorCredential(ctx context.Context, project identity.Principal, environment, key string, rotate bool) (IssuedExecutorCredential, error) { + target, err := s.selfHostedTarget(ctx, project, environment) + if err != nil { + return IssuedExecutorCredential{}, err + } + if err := s.checkActiveProject(ctx, project.TenantID); err != nil { + return IssuedExecutorCredential{}, err + } + if !rotate { + return s.issueExecutorCredential(ctx, project, key, target.ID, auditActiveProject("issue", key)) + } + if err := s.checkRestrictedTo(ctx, project, target, key); err != nil { + return IssuedExecutorCredential{}, err + } + return s.rotateExecutorCredential(ctx, project, key, auditActiveProject("rotate", key)) +} + +// RevokeProjectExecutorCredential revokes a key restricted to the +// Environment. Revoking it again succeeds, and so does revoking in an archived +// Project. +func (s *Service) RevokeProjectExecutorCredential(ctx context.Context, project identity.Principal, environment, key string) error { + target, err := s.selfHostedTarget(ctx, project, environment) + if err != nil { + return err + } + if err := s.checkRestrictedTo(ctx, project, target, key); err != nil { + return err + } + return s.storage.WithExecutorCredentials(ctx, project.TenantID, func(ctx context.Context, tx ExecutorCredentialTx) error { + if err := tx.RevokeExecutorCredential(ctx, project.Subject(), key); err != nil { + return err + } + return tx.RecordExecutorCredentialAudit(ctx, "revoke", key) + }) +} + +// auditActiveProject repeats the archive check in the writing transaction, +// whose share lock on the Project makes an issuance or rotation either commit +// before an archive or see it, then records the audit entry. +func auditActiveProject(action, key string) func(context.Context, ExecutorCredentialTx) error { + return func(ctx context.Context, tx ExecutorCredentialTx) error { + if err := lockActiveProject(ctx, tx); err != nil { + return err + } + return tx.RecordExecutorCredentialAudit(ctx, action, key) + } +} + +// checkRestrictedTo confirms that the Project's key is restricted to exactly +// the target; any other key is ErrNotFound. Restrictions and principals never +// change, so checking before the rotation or revocation transaction cannot +// authorize another target. +func (s *Service) checkRestrictedTo(ctx context.Context, project identity.Principal, target Environment, key string) error { + if err := checkExecutorIdentity(project, key); err != nil { + return err + } + restriction, err := s.storage.LoadExecutorCredentialRestriction(ctx, project, key) + if err != nil { + return err + } + if restriction != target.ID { + return ErrNotFound + } + return nil +} diff --git a/services/core/internal/sessions/service.go b/services/core/internal/sessions/service.go index d75300cbb..4ac82fdc9 100644 --- a/services/core/internal/sessions/service.go +++ b/services/core/internal/sessions/service.go @@ -18,4 +18,5 @@ func NewService(storage Storage) (*Service, error) { type Storage interface { ArtifactStorage DeviceStorage + ExecutorCredentialStorage } diff --git a/services/core/internal/sessions/transaction_test.go b/services/core/internal/sessions/transaction_test.go index 3cb35c1e8..2f54fb6f2 100644 --- a/services/core/internal/sessions/transaction_test.go +++ b/services/core/internal/sessions/transaction_test.go @@ -10,6 +10,7 @@ import ( "time" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/items" ) @@ -51,6 +52,20 @@ type fakeTx struct { bindDevice func() error authorizeEnrollment func() (EnrollmentAuthority, error) enrollDevice func() (string, error) + loadFileWrite func() (EnvironmentFileWrite, bool, error) + loadPendingInput func() (bool, error) + createFileWrite func() (EnvironmentFileWrite, error) + settleFileWrite func() (EnvironmentFileWrite, error) + recordFileWriteAudit func() error + + lockProject func() (bool, error) + listExecutorCredentials func() ([]ExecutorCredential, error) + authenticateExecutor func() (bool, error) + issueExecutorCredential func(ExecutorCredentialGrant) (IssuedExecutorCredential, error) + rotateExecutorCredential func(digest string) (IssuedExecutorCredential, error) + revokeExecutorCredential func() error + recordExecutorCredentialAudit func() error + loadSessionCreator func() (identity.Subject, bool, error) loadJournalTurn func() (JournalTurn, bool, error) matchEvents func() (bool, error) @@ -88,6 +103,10 @@ var ( _ ConnectionTx = (*fakeTx)(nil) _ DeviceBindingTx = (*fakeTx)(nil) _ EnrollmentTx = (*fakeTx)(nil) + _ FileWriteReservationTx = (*fakeTx)(nil) + _ FileWriteSettlementTx = (*fakeTx)(nil) + + _ EnvironmentExecutorCredentialTx = (*fakeTx)(nil) ) func (f *fakeTx) record(name string, set bool, detail ...string) { diff --git a/services/core/internal/store/device_bootstrap_binding_test.go b/services/core/internal/store/device_bootstrap_binding_test.go index d61c74381..d0c94f706 100644 --- a/services/core/internal/store/device_bootstrap_binding_test.go +++ b/services/core/internal/store/device_bootstrap_binding_test.go @@ -82,7 +82,7 @@ func TestDeviceCredentialWithoutManagedNodeRetainsPublicRouteIdentity(t *testing t.Fatalf("non-node credential acquired allocation route: found=%v node=%s error=%v", found, credential.RuntimeNodeID, err) } } - if err := s.RevokeExecutorCredential(t.Context(), principal, key.KeyID); err != nil { + if err := sessionService(t, s).RevokeExecutorCredential(t.Context(), principal, key.KeyID); err != nil { t.Fatal(err) } if _, found, err := sessionAdapter(s).GetDeviceCredential(t.Context(), enrolled.DeviceID); err != nil || found { diff --git a/services/core/internal/store/environment_executor_credentials.go b/services/core/internal/store/environment_executor_credentials.go deleted file mode 100644 index d9726abd2..000000000 --- a/services/core/internal/store/environment_executor_credentials.go +++ /dev/null @@ -1,136 +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/identity" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -// IssueExecutorCredential returns a new connect-only secret once, without replacing an existing ID. -func (s *Store) IssueExecutorCredential(ctx context.Context, principal identity.Principal, keyID, environment string) (sessions.IssuedExecutorCredential, error) { - return s.issueExecutorCredential(ctx, principal, keyID, environment, nil) -} - -// issueExecutorCredential runs before, when given, first in the issuing -// transaction; its error aborts the issuance. -func (s *Store) issueExecutorCredential(ctx context.Context, principal identity.Principal, keyID, environment string, before func(context.Context, *sqlc.Queries) error) (sessions.IssuedExecutorCredential, error) { - tenant, id, err := executorCredentialIdentity(principal, keyID) - if err != nil { - return sessions.IssuedExecutorCredential{}, err - } - exists, err := s.queries.ExecutorProjectScopeExists(ctx, sqlc.ExecutorProjectScopeExistsParams{TenantID: tenant, OrganizationID: principal.OrganizationID, ProjectID: principal.ProjectID}) - if err != nil { - return sessions.IssuedExecutorCredential{}, err - } - if !exists { - return sessions.IssuedExecutorCredential{}, sessions.ErrNotFound - } - var restriction pgtype.UUID - if environment != "" { - restriction, err = parseID(environment) - if err != nil { - return sessions.IssuedExecutorCredential{}, err - } - } - token, digest, err := newExecutorSecret() - if err != nil { - return sessions.IssuedExecutorCredential{}, err - } - var result sessions.IssuedExecutorCredential - err = s.withExecutorCredentialTarget(ctx, principal, restriction, func(ctx context.Context, q *sqlc.Queries) error { - if before != nil { - if err := before(ctx, q); err != nil { - return err - } - } - row, err := q.IssueExecutorCredential(ctx, sqlc.IssueExecutorCredentialParams{ - KeyID: id, TenantID: tenant, SubjectKind: pgtype.Text{String: principal.SubjectKind, Valid: true}, SubjectID: pgtype.Text{String: principal.SubjectID, Valid: true}, - OrganizationID: principal.OrganizationID, ProjectID: principal.ProjectID, EnvironmentID: restriction, TokenSha256: digest, - }) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrExecutorCredentialExists - } - if err != nil { - return err - } - result = issuedExecutorCredential(row.KeyID, row.EnvironmentID, token) - return nil - }) - if err != nil { - return sessions.IssuedExecutorCredential{}, err - } - return result, nil -} - -func (s *Store) RotateExecutorCredential(ctx context.Context, principal identity.Principal, keyID string) (sessions.IssuedExecutorCredential, error) { - return s.rotateExecutorCredential(ctx, principal, keyID, nil) -} - -// rotateExecutorCredential runs before, when given, first in the rotating -// transaction; its error aborts the rotation. -func (s *Store) rotateExecutorCredential(ctx context.Context, principal identity.Principal, keyID string, before func(context.Context, *sqlc.Queries) error) (sessions.IssuedExecutorCredential, error) { - tenant, id, err := executorCredentialIdentity(principal, keyID) - if err != nil { - return sessions.IssuedExecutorCredential{}, err - } - restriction, err := s.executorCredentialRestriction(ctx, principal, tenant, id) - if err != nil { - return sessions.IssuedExecutorCredential{}, err - } - token, digest, err := newExecutorSecret() - if err != nil { - return sessions.IssuedExecutorCredential{}, err - } - var result sessions.IssuedExecutorCredential - err = s.withExecutorCredentialTarget(ctx, principal, restriction, func(ctx context.Context, q *sqlc.Queries) error { - if before != nil { - if err := before(ctx, q); err != nil { - return err - } - } - row, err := q.RotateExecutorCredential(ctx, sqlc.RotateExecutorCredentialParams{KeyID: id, TenantID: tenant, SubjectKind: pgtype.Text{String: principal.SubjectKind, Valid: true}, SubjectID: pgtype.Text{String: principal.SubjectID, Valid: true}, TokenSha256: digest}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - result = issuedExecutorCredential(row.KeyID, row.EnvironmentID, token) - return nil - }) - if err != nil { - return sessions.IssuedExecutorCredential{}, err - } - return result, nil -} - -func (s *Store) RevokeExecutorCredential(ctx context.Context, principal identity.Principal, keyID string) error { - tenant, id, err := executorCredentialIdentity(principal, keyID) - if err != nil { - return err - } - if _, err := s.executorCredentialRestriction(ctx, principal, tenant, id); err != nil { - return err - } - n, err := s.queries.RevokeExecutorCredential(ctx, sqlc.RevokeExecutorCredentialParams{KeyID: id, TenantID: tenant, SubjectKind: pgtype.Text{String: principal.SubjectKind, Valid: true}, SubjectID: pgtype.Text{String: principal.SubjectID, Valid: true}}) - if err == nil && n == 0 { - return sessions.ErrNotFound - } - return err -} - -func (s *Store) executorCredentialRestriction(ctx context.Context, principal identity.Principal, tenant, id pgtype.UUID) (pgtype.UUID, error) { - restriction, err := s.queries.GetExecutorCredentialForPrincipal(ctx, sqlc.GetExecutorCredentialForPrincipalParams{ - KeyID: id, TenantID: tenant, SubjectKind: pgtype.Text{String: principal.SubjectKind, Valid: true}, SubjectID: pgtype.Text{String: principal.SubjectID, Valid: true}, - OrganizationID: principal.OrganizationID, ProjectID: principal.ProjectID, - }) - if errors.Is(err, pgx.ErrNoRows) { - return pgtype.UUID{}, sessions.ErrNotFound - } - return restriction, err -} diff --git a/services/core/internal/store/environment_executor_credentials_test.go b/services/core/internal/store/environment_executor_credentials_test.go index af9ad84fa..033ec501f 100644 --- a/services/core/internal/store/environment_executor_credentials_test.go +++ b/services/core/internal/store/environment_executor_credentials_test.go @@ -1,7 +1,9 @@ package store import ( + "crypto/rand" "crypto/sha256" + "encoding/base64" "encoding/hex" "errors" "sync" @@ -17,11 +19,22 @@ func executorDigest(token string) string { return hex.EncodeToString(sum[:]) } +// newExecutorSecret returns a secret an installer generates and its digest. +func newExecutorSecret() (string, string, error) { + secret := make([]byte, 32) + if _, err := rand.Read(secret); err != nil { + return "", "", err + } + token := base64.RawURLEncoding.EncodeToString(secret) + return token, executorDigest(token), nil +} + func TestEnvironmentExecutorCredentialLifecycle(t *testing.T) { s, pool := testStore(t) tenant, foreign := uuid.NewString(), uuid.NewString() ctx := t.Context() principal := FixtureExecutorPrincipal(t, s, tenant) + credentials := sessionService(t, s) session, err := s.CreateSession(ctx, tenant, environmentInput("credential", "self_hosted", "/workspace")) if err != nil { t.Fatal(err) @@ -30,13 +43,13 @@ func TestEnvironmentExecutorCredentialLifecycle(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.IssueExecutorCredential(ctx, FixtureExecutorPrincipal(t, s, foreign), environment.ID, environment.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := credentials.IssueExecutorCredential(ctx, FixtureExecutorPrincipal(t, s, foreign), environment.ID, environment.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign issue", err) } - if _, err := s.RotateExecutorCredential(ctx, principal, environment.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := credentials.RotateExecutorCredential(ctx, principal, environment.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("rotate manufactured a credential", err) } - issued, err := s.IssueExecutorCredential(ctx, principal, environment.ID, environment.ID) + issued, err := credentials.IssueExecutorCredential(ctx, principal, environment.ID, environment.ID) token := issued.Token if err != nil || len(token) != 43 { t.Fatal("issue failed", err) @@ -57,13 +70,13 @@ func TestEnvironmentExecutorCredentialLifecycle(t *testing.T) { 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) { + if _, err := credentials.IssueExecutorCredential(ctx, principal, environment.ID, environment.ID); !errors.Is(err, sessions.ErrExecutorCredentialExists) { t.Fatal("issue silently replaced credential", err) } - if err := s.RevokeExecutorCredential(ctx, FixtureExecutorPrincipal(t, s, foreign), environment.ID); !errors.Is(err, sessions.ErrNotFound) { + if err := credentials.RevokeExecutorCredential(ctx, FixtureExecutorPrincipal(t, s, foreign), environment.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign revoke", err) } - if _, err := s.RotateExecutorCredential(ctx, FixtureExecutorPrincipal(t, s, foreign), environment.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := credentials.RotateExecutorCredential(ctx, FixtureExecutorPrincipal(t, s, foreign), environment.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign rotate", err) } var stored string @@ -76,22 +89,22 @@ func TestEnvironmentExecutorCredentialLifecycle(t *testing.T) { } restarted, _ := testStore(t) check(restarted, token, true) - next, err := restarted.RotateExecutorCredential(ctx, principal, environment.ID) + next, err := sessionService(t, restarted).RotateExecutorCredential(ctx, principal, environment.ID) if err != nil || next.Token == token { t.Fatal("rotation failed", err) } check(s, token, false) check(s, next.Token, true) for range 2 { - if err := s.RevokeExecutorCredential(ctx, principal, environment.ID); err != nil { + if err := credentials.RevokeExecutorCredential(ctx, principal, environment.ID); err != nil { t.Fatal(err) } } check(restarted, next.Token, false) - if _, err := s.IssueExecutorCredential(ctx, principal, environment.ID, environment.ID); !errors.Is(err, sessions.ErrExecutorCredentialExists) { + if _, err := credentials.IssueExecutorCredential(ctx, principal, environment.ID, environment.ID); !errors.Is(err, sessions.ErrExecutorCredentialExists) { t.Fatal("ordinary issue resurrected revoked authority", err) } - restored, err := s.RotateExecutorCredential(ctx, principal, environment.ID) + restored, err := credentials.RotateExecutorCredential(ctx, principal, environment.ID) if err != nil { t.Fatal(err) } @@ -101,7 +114,7 @@ func TestEnvironmentExecutorCredentialLifecycle(t *testing.T) { t.Fatal(err) } check(restarted, restored.Token, false) - if _, err := s.RotateExecutorCredential(ctx, principal, environment.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := credentials.RotateExecutorCredential(ctx, principal, environment.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted Session authority resurrected", err) } if err := pool.QueryRow(ctx, "SELECT token_sha256 FROM environment_executor_credentials WHERE environment_id=$1", environment.ID).Scan(&stored); err != nil || stored != executorDigest(restored.Token) { @@ -125,6 +138,7 @@ func TestEnvironmentExecutorConcurrentIssueAndDeletion(t *testing.T) { } // Provisioning remains control-plane work while the execution owner is active. executionWriter(t, s) + services := []*sessions.Service{sessionService(t, s), sessionService(t, other)} const attempts = 8 var wg sync.WaitGroup tokens := make(chan string, attempts) @@ -132,11 +146,8 @@ func TestEnvironmentExecutorConcurrentIssueAndDeletion(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - st := s - if i%2 != 0 { - st = other - } - token, err := st.IssueExecutorCredential(ctx, principal, environment.ID, environment.ID) + service := services[i%2] + token, err := service.IssueExecutorCredential(ctx, principal, environment.ID, environment.ID) if err == nil { tokens <- token.Token } else if !errors.Is(err, sessions.ErrExecutorCredentialExists) { @@ -154,7 +165,7 @@ func TestEnvironmentExecutorConcurrentIssueAndDeletion(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - token, err := other.RotateExecutorCredential(ctx, principal, environment.ID) + token, err := services[1].RotateExecutorCredential(ctx, principal, environment.ID) if err == nil { rotated <- token.Token } else if !errors.Is(err, sessions.ErrNotFound) { diff --git a/services/core/internal/store/environment_executor_management.go b/services/core/internal/store/environment_executor_management.go deleted file mode 100644 index 361915520..000000000 --- a/services/core/internal/store/environment_executor_management.go +++ /dev/null @@ -1,210 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - - "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/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" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -// The Project* methods serve Core-key executor credential management. project -// is the Project's own principal, which becomes the credential's execution -// principal; the target must be a self_hosted Environment of that Project whose -// Session is not deleted, and only credentials restricted to it are managed. - -func (s *Store) ListProjectExecutorCredentials(ctx context.Context, project identity.Principal, environment string) ([]sessions.ExecutorCredential, error) { - state, err := s.ProjectExecutorCredentialState(ctx, project, environment) - return state.Credentials, err -} - -// ProjectExecutorCredentialState reads list metadata and binding facts in one -// read-only snapshot. The snapshot ends before any live peer/authority observation. -func (s *Store) ProjectExecutorCredentialState(ctx context.Context, project identity.Principal, environment string) (sessions.ExecutorCredentialState, error) { - if err := project.Validate(); err != nil { - return sessions.ExecutorCredentialState{}, sessions.ErrInvalidInput - } - tenant, err := parseID(project.TenantID) - if err != nil { - return sessions.ExecutorCredentialState{}, err - } - environmentID := pgunit.PathID(environment) - result := sessions.ExecutorCredentialState{Credentials: []sessions.ExecutorCredential{}} - err = s.pooled.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - row, err := q.GetEnvironmentExecutorConnection(ctx, sqlc.GetEnvironmentExecutorConnectionParams{EnvironmentID: environmentID, TenantID: tenant}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - result.EnvironmentID = uuid.UUID(environmentID.Bytes).String() - result.Connection.EnvironmentStatus = row.EnvironmentStatus - if row.DeviceID.Valid { - result.Connection.DeviceID = uuid.UUID(row.DeviceID.Bytes).String() - } - if row.ExecutorKeyID.Valid { - key := uuid.UUID(row.ExecutorKeyID.Bytes).String() - result.Connection.BoundKeyID = &key - } - if row.EnrolledAt.Valid { - at := row.EnrolledAt.Time - result.Connection.EnrolledAt = &at - } - if row.LastSeenAt.Valid { - at := row.LastSeenAt.Time - result.Connection.LastSeenAt = &at - } - if row.CredentialHash.Valid { - result.Connection.CredentialHash = row.CredentialHash.String - } - rows, err := q.ListEnvironmentExecutorCredentials(ctx, sqlc.ListEnvironmentExecutorCredentialsParams{ - TenantID: tenant, EnvironmentID: environmentID, - SubjectKind: pgtype.Text{String: project.SubjectKind, Valid: true}, SubjectID: pgtype.Text{String: project.SubjectID, Valid: true}, - }) - if err != nil { - return err - } - for _, row := range rows { - credential := sessions.ExecutorCredential{KeyID: uuid.UUID(row.KeyID.Bytes).String(), CreatedAt: row.CreatedAt.Time} - if row.RevokedAt.Valid { - revoked := row.RevokedAt.Time - credential.RevokedAt = &revoked - } - result.Credentials = append(result.Credentials, credential) - } - return nil - }) - return result, err -} - -// IssueProjectExecutorCredential issues a new key or, with rotate, replaces the -// secret of an existing key restricted to the Environment. An archived Project -// gets neither (projects.ErrArchived). The administrator audit entry commits in -// the same transaction and never contains the secret. -func (s *Store) IssueProjectExecutorCredential(ctx context.Context, project identity.Principal, environment, keyID string, rotate bool) (sessions.IssuedExecutorCredential, error) { - // The target is checked first, then the archived Project, then the key. - if err := s.selfHostedExecutorTarget(ctx, project, environment); err != nil { - return sessions.IssuedExecutorCredential{}, err - } - if err := activeProject(ctx, s.queries, project); err != nil { - return sessions.IssuedExecutorCredential{}, err - } - if !rotate { - return s.issueExecutorCredential(ctx, project, keyID, environment, activeProjectAudit(project, "issue", keyID)) - } - if err := s.exactExecutorRestriction(ctx, project, environment, keyID); err != nil { - return sessions.IssuedExecutorCredential{}, err - } - return s.rotateExecutorCredential(ctx, project, keyID, activeProjectAudit(project, "rotate", keyID)) -} - -// RevokeProjectExecutorCredential is idempotent and also works in an archived -// Project; each successful request is audited. -func (s *Store) RevokeProjectExecutorCredential(ctx context.Context, project identity.Principal, environment, keyID string) error { - if err := s.selfHostedExecutorTarget(ctx, project, environment); err != nil { - return err - } - if err := s.exactExecutorRestriction(ctx, project, environment, keyID); err != nil { - return err - } - tenant, id, err := executorCredentialIdentity(project, keyID) - if err != nil { - return err - } - record := executorCredentialAudit(project, "revoke", keyID) - return s.pooled.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - n, err := q.RevokeExecutorCredential(ctx, sqlc.RevokeExecutorCredentialParams{KeyID: id, TenantID: tenant, SubjectKind: pgtype.Text{String: project.SubjectKind, Valid: true}, SubjectID: pgtype.Text{String: project.SubjectID, Valid: true}}) - if err != nil { - return err - } - if n == 0 { - return sessions.ErrNotFound - } - return record(ctx, q) - }) -} - -func executorCredentialAudit(project identity.Principal, action, keyID string) func(context.Context, *sqlc.Queries) error { - return func(ctx context.Context, q *sqlc.Queries) error { - return auditpg.RecordAdminMutation(ctx, q, project.TenantID, action, "executor_credential", keyID) - } -} - -// activeProjectAudit repeats the archive check in the writing transaction. It -// share-locks the Project, which archiving updates, so an issuance or rotation -// either commits before the archive or sees it and fails. -func activeProjectAudit(project identity.Principal, action, keyID string) func(context.Context, *sqlc.Queries) error { - audit := executorCredentialAudit(project, action, keyID) - return func(ctx context.Context, q *sqlc.Queries) error { - if err := activeProject(ctx, q, project); err != nil { - return err - } - return audit(ctx, q) - } -} - -// activeProject returns projects.ErrArchived for an archived Project. -func activeProject(ctx context.Context, q *sqlc.Queries, project identity.Principal) error { - tenant, err := parseID(project.TenantID) - if err != nil { - return err - } - row, err := q.LockProjectByTenant(ctx, tenant) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - if row.ArchivedAt.Valid { - return projects.ErrArchived - } - return nil -} - -func (s *Store) selfHostedExecutorTarget(ctx context.Context, principal identity.Principal, environment string) error { - if err := principal.Validate(); err != nil { - return sessions.ErrInvalidInput - } - owned, err := sessionpg.LoadEnvironment(ctx, s.queries, principal.TenantID, environment) - if err != nil { - return err - } - var configuration struct { - Type string `json:"type"` - } - if json.Unmarshal(owned.Configuration, &configuration) != nil || configuration.Type != "self_hosted" { - return sessions.ErrNotFound - } - return nil -} - -func (s *Store) exactExecutorRestriction(ctx context.Context, principal identity.Principal, environment, keyID string) error { - tenant, id, err := executorCredentialIdentity(principal, keyID) - if err != nil { - return err - } - want := pgunit.PathID(environment) - actual, err := s.executorCredentialRestriction(ctx, principal, tenant, id) - if err != nil { - return err - } - // Restrictions and principals are immutable, so checking before the - // rotation or revocation transaction cannot authorize a different target. - if !actual.Valid || actual != want { - return sessions.ErrNotFound - } - return nil -} diff --git a/services/core/internal/store/environment_executor_management_test.go b/services/core/internal/store/environment_executor_management_test.go index 2e87de334..3e0b373d8 100644 --- a/services/core/internal/store/environment_executor_management_test.go +++ b/services/core/internal/store/environment_executor_management_test.go @@ -8,11 +8,19 @@ import ( "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/adminaudit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" ) +// projectCredentials lists the credentials of the Project's principal +// restricted to the Environment. +func projectCredentials(ctx context.Context, s *Store, project identity.Principal, environment string) ([]sessions.ExecutorCredential, error) { + state, err := sessionAdapter(s).ProjectExecutorCredentialState(ctx, project, environment) + return state.Credentials, err +} + func TestProjectEnvironmentExecutorManagement(t *testing.T) { s, pool := testStore(t) ctx := t.Context() @@ -42,10 +50,10 @@ func TestProjectEnvironmentExecutorManagement(t *testing.T) { keyID := uuid.NewString() // The audit entry commits with the write: without an audit source nothing is issued. - if _, err := s.IssueProjectExecutorCredential(ctx, p, one.ID, keyID, false); !errors.Is(err, adminaudit.ErrInvalidSource) { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(ctx, p, one.ID, keyID, false); !errors.Is(err, adminaudit.ErrInvalidSource) { t.Fatal("unaudited issue", err) } - issued, err := s.IssueProjectExecutorCredential(admin(), p, one.ID, keyID, false) + issued, err := sessionService(t, s).IssueProjectExecutorCredential(admin(), p, one.ID, keyID, false) if err != nil || issued.KeyID != keyID || issued.EnvironmentID != one.ID || issued.Token == "" { t.Fatal("issue", err) } @@ -53,40 +61,40 @@ func TestProjectEnvironmentExecutorManagement(t *testing.T) { t.Fatal("issued credential does not authenticate", err) } // A lost issuance response must not cause a new key or replace the old secret. - if _, err := s.IssueProjectExecutorCredential(admin(), p, one.ID, keyID, false); !errors.Is(err, sessions.ErrExecutorCredentialExists) { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(admin(), p, one.ID, keyID, false); !errors.Is(err, sessions.ErrExecutorCredentialExists) { t.Fatal("uncertain retry", err) } - listed, err := s.ListProjectExecutorCredentials(ctx, p, one.ID) + listed, err := projectCredentials(ctx, s, p, one.ID) if err != nil || len(listed) != 1 || listed[0].KeyID != keyID || listed[0].CreatedAt.IsZero() || listed[0].RevokedAt != nil { t.Fatal("list", listed, err) } - if listed, err := s.ListProjectExecutorCredentials(ctx, p, two.ID); err != nil || len(listed) != 0 { + if listed, err := projectCredentials(ctx, s, p, two.ID); err != nil || len(listed) != 0 { t.Fatal("other Environment list", listed, err) } // Another Project, a hosted or unknown Environment and a key restricted elsewhere are not found. - if _, err := s.ListProjectExecutorCredentials(ctx, foreign.Principal, one.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := projectCredentials(ctx, s, foreign.Principal, one.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign list", err) } - if _, err := s.IssueProjectExecutorCredential(keyAdminContext(ctx, foreignProject.ID), foreign.Principal, one.ID, uuid.NewString(), false); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(keyAdminContext(ctx, foreignProject.ID), foreign.Principal, one.ID, uuid.NewString(), false); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign issue", err) } for _, environment := range []string{two.ID, hosted.ID, uuid.NewString()} { - if _, err := s.IssueProjectExecutorCredential(admin(), p, environment, keyID, true); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(admin(), p, environment, keyID, true); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("wrong target rotate", err) } - if err := s.RevokeProjectExecutorCredential(admin(), p, environment, keyID); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionService(t, s).RevokeProjectExecutorCredential(admin(), p, environment, keyID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("wrong target revoke", err) } } - if _, err := s.IssueProjectExecutorCredential(admin(), p, hosted.ID, uuid.NewString(), false); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(admin(), p, hosted.ID, uuid.NewString(), false); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("hosted issuance", err) } - if _, err := s.IssueProjectExecutorCredential(admin(), p, one.ID, uuid.NewString(), true); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(admin(), p, one.ID, uuid.NewString(), true); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unknown key rotation", err) } - rotated, err := s.IssueProjectExecutorCredential(admin(), p, one.ID, keyID, true) + rotated, err := sessionService(t, s).IssueProjectExecutorCredential(admin(), p, one.ID, keyID, true) if err != nil || rotated.Token == issued.Token { t.Fatal("rotation", err) } @@ -94,17 +102,17 @@ func TestProjectEnvironmentExecutorManagement(t *testing.T) { t.Fatal("old secret", err) } for range 2 { - if err := s.RevokeProjectExecutorCredential(admin(), p, one.ID, keyID); err != nil { + if err := sessionService(t, s).RevokeProjectExecutorCredential(admin(), p, one.ID, keyID); err != nil { t.Fatal(err) } } 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 { + if listed, err := projectCredentials(ctx, s, p, one.ID); err != nil || len(listed) != 1 || listed[0].RevokedAt == nil { t.Fatal("revoked list", listed, err) } - if _, err := s.IssueProjectExecutorCredential(admin(), p, one.ID, keyID, false); !errors.Is(err, sessions.ErrExecutorCredentialExists) { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(admin(), p, one.ID, keyID, false); !errors.Is(err, sessions.ErrExecutorCredentialExists) { t.Fatal("resurrected secret", err) } @@ -130,13 +138,13 @@ func TestProjectEnvironmentExecutorManagement(t *testing.T) { if err := s.DeleteSession(ctx, p.TenantID, session.ID); err != nil { t.Fatal(err) } - if _, err := s.ListProjectExecutorCredentials(ctx, p, one.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := projectCredentials(ctx, s, p, one.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted target list", err) } - if _, err := s.IssueProjectExecutorCredential(admin(), p, one.ID, keyID, true); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(admin(), p, one.ID, keyID, true); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted target rotate", err) } - if err := s.RevokeProjectExecutorCredential(admin(), p, one.ID, keyID); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionService(t, s).RevokeProjectExecutorCredential(admin(), p, one.ID, keyID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted target revoke", err) } } @@ -158,7 +166,7 @@ func TestArchivedProjectExecutorCredentials(t *testing.T) { t.Fatal(err) } keyID := uuid.NewString() - if _, err := s.IssueProjectExecutorCredential(keyAdminContext(ctx, project.ID), p, environment.ID, keyID, false); err != nil { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(keyAdminContext(ctx, project.ID), p, environment.ID, keyID, false); err != nil { t.Fatal(err) } if _, err := testProjects(t, pool).ArchiveProject(keyAdminContext(ctx, project.ID), projects.ArchiveProject{ID: project.ID}); err != nil { @@ -177,17 +185,17 @@ func TestArchivedProjectExecutorCredentials(t *testing.T) { {environment.ID, keyID, true, projects.ErrArchived}, {environment.ID, uuid.NewString(), true, projects.ErrArchived}, } { - if _, err := s.IssueProjectExecutorCredential(keyAdminContext(ctx, project.ID), p, test.environment, test.key, test.rotate); !errors.Is(err, test.want) { + if _, err := sessionService(t, s).IssueProjectExecutorCredential(keyAdminContext(ctx, project.ID), p, test.environment, test.key, test.rotate); !errors.Is(err, test.want) { t.Fatal("archived write", test, err) } } - if listed, err := s.ListProjectExecutorCredentials(ctx, p, environment.ID); err != nil || len(listed) != 1 { + if listed, err := projectCredentials(ctx, s, p, environment.ID); err != nil || len(listed) != 1 { t.Fatal("archived list", listed, err) } - if err := s.RevokeProjectExecutorCredential(keyAdminContext(ctx, project.ID), p, environment.ID, keyID); err != nil { + if err := sessionService(t, s).RevokeProjectExecutorCredential(keyAdminContext(ctx, project.ID), p, environment.ID, keyID); err != nil { t.Fatal("archived revoke", err) } - if listed, err := s.ListProjectExecutorCredentials(ctx, p, environment.ID); err != nil || len(listed) != 1 || listed[0].RevokedAt == nil { + if listed, err := projectCredentials(ctx, s, p, environment.ID); err != nil || len(listed) != 1 || listed[0].RevokedAt == nil { t.Fatal("revoked list", listed, err) } } @@ -198,7 +206,7 @@ func TestProjectExecutorConnectionState(t *testing.T) { binding := createTestProject(t, pool) project := binding.Project session, env, key := runtimeEnrollmentFixture(t, s, binding.Principal) - state, err := s.ProjectExecutorCredentialState(ctx, binding.Principal, env.ID) + state, err := sessionAdapter(s).ProjectExecutorCredentialState(ctx, binding.Principal, env.ID) if err != nil || state.Connection.DeviceID != "" || state.Connection.EnrolledAt != nil || state.Connection.BoundKeyID != nil || state.Connection.LastSeenAt != nil { t.Fatal("never enrolled", state, err) } @@ -208,7 +216,7 @@ func TestProjectExecutorConnectionState(t *testing.T) { } check := func() sessions.ExecutorCredentialState { t.Helper() - v, e := s.ProjectExecutorCredentialState(ctx, binding.Principal, env.ID) + v, e := sessionAdapter(s).ProjectExecutorCredentialState(ctx, binding.Principal, env.ID) if e != nil { t.Fatal(e) } @@ -219,13 +227,13 @@ func TestProjectExecutorConnectionState(t *testing.T) { t.Fatal("binding", state) } for _, spelling := range []string{env.ID, strings.ToUpper(env.ID), strings.ReplaceAll(env.ID, "-", "")} { - resolved, err := s.ProjectExecutorCredentialState(ctx, binding.Principal, spelling) + resolved, err := sessionAdapter(s).ProjectExecutorCredentialState(ctx, binding.Principal, spelling) if err != nil || resolved.EnvironmentID != env.ID || resolved.Connection.DeviceID != enrolled.DeviceID || resolved.Connection.CredentialHash != executorDigest(key.Token) { t.Fatal("equivalent target did not retain canonical identity and binding", spelling, resolved, err) } } // An additional credential never changes the enrolled device's bound key. - if _, err = s.IssueProjectExecutorCredential(keyAdminContext(ctx, project.ID), binding.Principal, env.ID, uuid.NewString(), false); err != nil { + if _, err = sessionService(t, s).IssueProjectExecutorCredential(keyAdminContext(ctx, project.ID), binding.Principal, env.ID, uuid.NewString(), false); err != nil { t.Fatal(err) } state = check() @@ -239,7 +247,7 @@ func TestProjectExecutorConnectionState(t *testing.T) { if state.Connection.LastSeenAt == nil { t.Fatal("heartbeat history missing") } - rotated, err := s.IssueProjectExecutorCredential(keyAdminContext(ctx, project.ID), binding.Principal, env.ID, key.KeyID, true) + rotated, err := sessionService(t, s).IssueProjectExecutorCredential(keyAdminContext(ctx, project.ID), binding.Principal, env.ID, key.KeyID, true) if err != nil { t.Fatal(err) } @@ -260,7 +268,7 @@ func TestProjectExecutorConnectionState(t *testing.T) { if _, err = pool.Exec(ctx, "UPDATE environments SET status='connected' WHERE id=$1", env.ID); err != nil { t.Fatal(err) } - if err = s.RevokeProjectExecutorCredential(keyAdminContext(ctx, project.ID), binding.Principal, env.ID, key.KeyID); err != nil { + if err = sessionService(t, s).RevokeProjectExecutorCredential(keyAdminContext(ctx, project.ID), binding.Principal, env.ID, key.KeyID); err != nil { t.Fatal(err) } state = check() @@ -282,17 +290,17 @@ func TestProjectExecutorConnectionState(t *testing.T) { } foreign := createTestProject(t, pool) for _, id := range []string{uuid.NewString(), "malformed"} { - if _, err = s.ProjectExecutorCredentialState(ctx, binding.Principal, id); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).ProjectExecutorCredentialState(ctx, binding.Principal, id); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("missing target", err) } } - if _, err = s.ProjectExecutorCredentialState(ctx, foreign.Principal, env.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).ProjectExecutorCredentialState(ctx, foreign.Principal, env.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign target", err) } if _, err = pool.Exec(ctx, "UPDATE sessions SET deleted_at=clock_timestamp() WHERE id=$1", session.ID); err != nil { t.Fatal(err) } - if _, err = s.ProjectExecutorCredentialState(ctx, binding.Principal, env.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).ProjectExecutorCredentialState(ctx, binding.Principal, env.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted target", err) } } diff --git a/services/core/internal/store/environment_file_write_semantics_public_test.go b/services/core/internal/store/environment_file_write_semantics_public_test.go index 0de859659..ef032e80e 100644 --- a/services/core/internal/store/environment_file_write_semantics_public_test.go +++ b/services/core/internal/store/environment_file_write_semantics_public_test.go @@ -141,7 +141,7 @@ func TestEnvironmentFileCreateRejectionsLeaveNoReceiptOrConsumption(t *testing.T done := post(token, environment.ID, tc.body) id := serveFileWrite(h, rejected(tc.reason)) assertError(await(done), tc.message) - if intent, err := h.s.GetEnvironmentFileWrite(t.Context(), h.tenant, environment.ID, id); err != nil || intent.State != "rejected" { + if intent, err := store.FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, id); err != nil || intent.State != "rejected" { t.Fatal("rejection did not settle", intent, err) } } diff --git a/services/core/internal/store/environment_file_writes.go b/services/core/internal/store/environment_file_writes.go deleted file mode 100644 index f603cdfa9..000000000 --- a/services/core/internal/store/environment_file_writes.go +++ /dev/null @@ -1,194 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - - "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/auditpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/writeaudit" -) - -// ReserveEnvironmentFileWrite persists intent before external dispatch. A retry -// observes the earlier operation and never authorizes resending an unknown write. -func (s *Store) ReserveEnvironmentFileWrite(ctx context.Context, tenant, environment string, key sessions.FileWriteIdentity) (sessions.EnvironmentFileWrite, error) { - if err := s.checkExecutionAuthority(); err != nil { - return sessions.EnvironmentFileWrite{}, err - } - if !key.Valid() { - return sessions.EnvironmentFileWrite{}, sessions.ErrInvalidInput - } - owned, err := sessionpg.LoadEnvironment(ctx, s.queries, tenant, environment) - if err != nil { - return sessions.EnvironmentFileWrite{}, err - } - lookup, err := fileWriteLookup(tenant, environment, key.ID) - if err != nil { - return sessions.EnvironmentFileWrite{}, err - } - var origin []byte - if source, ok := writeaudit.FromContext(ctx); ok { - if err := source.Validate(tenant); err != nil { - return sessions.EnvironmentFileWrite{}, err - } - origin, err = json.Marshal(source) - if err != nil { - return sessions.EnvironmentFileWrite{}, err - } - } - var result sessions.EnvironmentFileWrite - err = s.withPublicSession(ctx, tenant, owned.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - previous, err := q.GetEnvironmentFileWrite(ctx, lookup) - if err == nil { - result = fileWriteFromRow(previous.EnvironmentFileWrite, session) - result.Replayed = true - if result.Identity != key { - return sessions.ErrIdempotencyConflict - } - return nil - } - if !errors.Is(err, pgx.ErrNoRows) { - return err - } - current, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: lookup.TenantID, ID: session}) - if err != nil { - return err - } - if current.Environment.ID != lookup.EnvironmentID || current.Environment.Status == "failed" || current.Environment.Status == "expired" { - return sessions.ErrInvalidInput - } - device, err := q.GetSessionDevice(ctx, sqlc.GetSessionDeviceParams{TenantID: lookup.TenantID, ID: session}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - if uuid.UUID(device.ID.Bytes).String() != key.DeviceID || device.EnvironmentID != lookup.EnvironmentID { - return sessions.ErrDeviceBindingConflict - } - bound := sessionpg.BindSession(q, lookup.TenantID, session) - if err := sessions.CheckComputeAdmission(ctx, bound); err != nil { - return err - } - if err := sessions.CheckInputStart(ctx, bound); err != nil { - return err - } - pending, err := q.EnvironmentFileWriteHasPendingInput(ctx, session) - if err != nil { - return err - } - if pending { - return sessions.ErrTurnConflict - } - deviceID, _ := parseID(key.DeviceID) - row, err := q.CreateEnvironmentFileWrite(ctx, sqlc.CreateEnvironmentFileWriteParams{ - ID: lookup.ID, EnvironmentID: lookup.EnvironmentID, DeviceID: deviceID, RequestSha256: key.RequestSHA256, AuditSource: origin, - }) - if err == nil { - result = fileWriteFromRow(row, session) - } - return err - }) - if err != nil { - return sessions.EnvironmentFileWrite{}, err - } - return result, nil -} - -// GetEnvironmentFileWrite also retains deleted-Session intents for internal -// cleanup. It is not a public resource query and never authorizes dispatch. -func (s *Store) GetEnvironmentFileWrite(ctx context.Context, tenant, environment, id string) (sessions.EnvironmentFileWrite, error) { - lookup, err := fileWriteLookup(tenant, environment, id) - if err != nil { - return sessions.EnvironmentFileWrite{}, err - } - row, err := s.queries.GetEnvironmentFileWrite(ctx, lookup) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.EnvironmentFileWrite{}, sessions.ErrNotFound - } - if err != nil { - return sessions.EnvironmentFileWrite{}, err - } - result := fileWriteFromRow(row.EnvironmentFileWrite, row.SessionID) - result.Replayed = true - return result, nil -} - -// SettleEnvironmentFileWrite requires an independently validated exact receipt. -// A missing receipt, cancellation or owner retirement is not a rejected upload. -func (s *Store) SettleEnvironmentFileWrite(ctx context.Context, tenant, environment string, key sessions.FileWriteIdentity, state string) (sessions.EnvironmentFileWrite, error) { - if err := s.checkExecutionAuthority(); err != nil { - return sessions.EnvironmentFileWrite{}, err - } - if !key.Valid() || (state != "committed" && state != "rejected") { - return sessions.EnvironmentFileWrite{}, sessions.ErrInvalidInput - } - previous, err := s.GetEnvironmentFileWrite(ctx, tenant, environment, key.ID) - if err != nil { - return sessions.EnvironmentFileWrite{}, err - } - lookup, _ := fileWriteLookup(tenant, environment, key.ID) - var result sessions.EnvironmentFileWrite - err = s.withSession(ctx, tenant, previous.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - current, err := q.GetEnvironmentFileWrite(ctx, lookup) - if err != nil { - return err - } - result = fileWriteFromRow(current.EnvironmentFileWrite, session) - if result.Identity != key || (result.State != "pending" && result.State != state) { - return sessions.ErrIdempotencyConflict - } - if result.State == state { - result.Replayed = true - return nil - } - row, err := q.SettleEnvironmentFileWrite(ctx, sqlc.SettleEnvironmentFileWriteParams{EnvironmentID: lookup.EnvironmentID, ID: lookup.ID, State: state}) - if err != nil { - return err - } - result = fileWriteFromRow(row, session) - if state == "committed" && len(row.AuditSource) > 0 { - var source writeaudit.Source - if err := json.Unmarshal(row.AuditSource, &source); err != nil { - return err - } - auditCtx := writeaudit.WithSource(ctx, source) - return auditpg.RecordWriteAudit(auditCtx, q, tenant, "upload_file", "environment", environment, uuid.UUID(session.Bytes).String()) - } - return nil - }) - if err != nil { - return sessions.EnvironmentFileWrite{}, err - } - return result, nil -} - -func fileWriteLookup(tenant, environment, id string) (sqlc.GetEnvironmentFileWriteParams, error) { - var result sqlc.GetEnvironmentFileWriteParams - var err error - result.TenantID, err = parseID(tenant) - if err == nil { - result.EnvironmentID, err = parseID(environment) - } - if err == nil { - result.ID, err = parseID(id) - } - return result, err -} - -func fileWriteFromRow(row sqlc.EnvironmentFileWrite, session pgtype.UUID) sessions.EnvironmentFileWrite { - result := sessions.EnvironmentFileWrite{Identity: sessions.FileWriteIdentity{ID: uuid.UUID(row.ID.Bytes).String(), DeviceID: uuid.UUID(row.DeviceID.Bytes).String(), RequestSHA256: row.RequestSha256}, - EnvironmentID: uuid.UUID(row.EnvironmentID.Bytes).String(), SessionID: uuid.UUID(session.Bytes).String(), State: row.State, CreatedAt: row.CreatedAt.Time} - if row.SettledAt.Valid { - result.SettledAt = &row.SettledAt.Time - } - return result -} diff --git a/services/core/internal/store/environment_file_writes_test.go b/services/core/internal/store/environment_file_writes_test.go index 51b5e7395..d4fb1b3a0 100644 --- a/services/core/internal/store/environment_file_writes_test.go +++ b/services/core/internal/store/environment_file_writes_test.go @@ -6,30 +6,33 @@ import ( "sync" "testing" + "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" "github.com/google/uuid" ) type fileWriteFixture struct { - s, writer *Store - tenant string - session sessions.Session - env sessions.Environment - key sessions.FileWriteIdentity + s *Store + lease *pgunit.Lease // the execution lease + writer *sessions.ExecutionOperations // runs the file writes on lease + tenant string + session sessions.Session + env sessions.Environment + key sessions.FileWriteIdentity } func newFileWriteFixture(t *testing.T) fileWriteFixture { t.Helper() s, pool := testStore(t) - writer := executionWriter(t, s) + lease := executionWriter(t, s).lease tenant := uuid.NewString() session, env := localEnvironment(t, s, tenant) host, err := FixtureEnvironmentDevice(t.Context(), pool, tenant, env.ID, "file owner", runtimedevice.HashCredential(uuid.NewString())) if err != nil { t.Fatal(err) } - return fileWriteFixture{s: s, writer: writer, tenant: tenant, session: session, env: env, + return fileWriteFixture{s: s, lease: lease, writer: sessionExecution(t, lease), tenant: tenant, session: session, env: env, key: sessions.FileWriteIdentity{ID: uuid.NewString(), DeviceID: host.ID, RequestSHA256: strings.Repeat("a", 64)}} } @@ -48,7 +51,7 @@ func TestEnvironmentFileWriteRetainsUnknownAcrossLeaseLoss(t *testing.T) { t.Fatal("lost writer settled an upload") } reopened, _ := testStore(t) - next := executionWriter(t, reopened) + next := sessionExecution(t, executionWriter(t, reopened).lease) got, err := next.ReserveEnvironmentFileWrite(ctx, f.tenant, f.env.ID, f.key) if err != nil || !got.Replayed || got.State != "pending" || !got.CreatedAt.Equal(first.CreatedAt) || got.Identity != f.key { t.Fatal("restart lost unknown write identity", got, err) @@ -73,7 +76,7 @@ func TestEnvironmentFileWriteRetainsUnknownAcrossLeaseLoss(t *testing.T) { if receipt, err := reopened.RequestCancel(ctx, f.tenant, f.session.ID, "idle-cancel"); err != nil || receipt.TurnID != "" { t.Fatal("write gate imposed mutation admission on idle cancellation", receipt, err) } - if got, err := reopened.GetEnvironmentFileWrite(ctx, f.tenant, f.env.ID, f.key.ID); err != nil || got.State != "pending" { + if got, err := FixtureFileWrite(ctx, reopened.pool, f.tenant, f.env.ID, f.key.ID); err != nil || got.State != "pending" { t.Fatal("idle cancellation cleared unknown write", got, err) } settled, err := next.SettleEnvironmentFileWrite(ctx, f.tenant, f.env.ID, f.key, "committed") @@ -86,9 +89,6 @@ func TestEnvironmentFileWriteRetainsUnknownAcrossLeaseLoss(t *testing.T) { func TestEnvironmentFileWriteMatchesReceiptAndRetainsDeletedOwner(t *testing.T) { f := newFileWriteFixture(t) ctx := t.Context() - if _, err := f.s.ReserveEnvironmentFileWrite(ctx, f.tenant, f.env.ID, f.key); err == nil { - t.Fatal("pooled Store admitted a file write") - } if _, err := f.writer.ReserveEnvironmentFileWrite(ctx, f.tenant, f.env.ID, f.key); err != nil { t.Fatal(err) } @@ -108,9 +108,6 @@ func TestEnvironmentFileWriteMatchesReceiptAndRetainsDeletedOwner(t *testing.T) if _, err := f.writer.SettleEnvironmentFileWrite(ctx, uuid.NewString(), f.env.ID, f.key, "committed"); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("cross-tenant settlement", err) } - if _, err := f.s.GetEnvironmentFileWrite(ctx, uuid.NewString(), f.env.ID, f.key.ID); !errors.Is(err, sessions.ErrNotFound) { - t.Fatal("cross-tenant read", err) - } for _, state := range []string{"pending", "unknown", "cancelled", "retired"} { if _, err := f.writer.SettleEnvironmentFileWrite(ctx, f.tenant, f.env.ID, f.key, state); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal("non-receipt settled write", state, err) @@ -119,7 +116,7 @@ func TestEnvironmentFileWriteMatchesReceiptAndRetainsDeletedOwner(t *testing.T) if err := f.s.DeleteSession(ctx, f.tenant, f.session.ID); err != nil { t.Fatal(err) } - if got, err := f.s.GetEnvironmentFileWrite(ctx, f.tenant, f.env.ID, f.key.ID); err != nil || got.State != "pending" { + if got, err := FixtureFileWrite(ctx, f.s.pool, f.tenant, f.env.ID, f.key.ID); err != nil || got.State != "pending" { t.Fatal("deletion discarded unresolved write", got, err) } if _, err := f.writer.ReserveEnvironmentFileWrite(ctx, f.tenant, f.env.ID, f.key); !errors.Is(err, sessions.ErrNotFound) { @@ -206,7 +203,7 @@ func TestEnvironmentFileWriteSerializesWithInputAndRetry(t *testing.T) { if got, err := f.writer.SettleEnvironmentFileWrite(ctx, f.tenant, f.env.ID, original, "committed"); err != nil || !got.Replayed { t.Fatal("old receipt retry changed", got, err) } - if got, err := f.s.GetEnvironmentFileWrite(ctx, f.tenant, f.env.ID, f.key.ID); err != nil || got.State != "pending" { + if got, err := FixtureFileWrite(ctx, f.s.pool, f.tenant, f.env.ID, f.key.ID); err != nil || got.State != "pending" { t.Fatal("old receipt cleared successor", got, err) } } diff --git a/services/core/internal/store/environment_initialization_test.go b/services/core/internal/store/environment_initialization_test.go index 66b8554f4..5504ef831 100644 --- a/services/core/internal/store/environment_initialization_test.go +++ b/services/core/internal/store/environment_initialization_test.go @@ -70,7 +70,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test if err != nil { t.Fatal(err) } - key, err := s.IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + key, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } @@ -187,7 +187,7 @@ func TestEnvironmentInitializationRevocationBeforeClaim(t *testing.T) { if err != nil { t.Fatal(err) } - key, err := s.IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + key, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_installation.go b/services/core/internal/store/environment_installation.go deleted file mode 100644 index ae3089053..000000000 --- a/services/core/internal/store/environment_installation.go +++ /dev/null @@ -1,108 +0,0 @@ -package store - -import ( - "context" - "crypto/hmac" - "crypto/sha256" - "encoding/base64" - "encoding/hex" - "encoding/json" - "errors" - "strings" - "time" - - "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/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -func (s *Store) AuthorizeEnvironmentInstallation(ctx context.Context, principal identity.Principal, environment, version string) (string, int64, error) { - if err := s.selfHostedExecutorTarget(ctx, principal, environment); err != nil { - return "", 0, err - } - if err := activeProject(ctx, s.queries, principal); err != nil { - return "", 0, err - } - claim := sessions.InstallationAuthorization{Principal: principal, Environment: environment, Version: version, ExpiresAt: time.Now().Add(30 * time.Minute).Unix()} - payload, err := json.Marshal(claim) - if err != nil { - return "", 0, err - } - encoded := base64.RawURLEncoding.EncodeToString(payload) - signature, err := s.credentialCipher.Fingerprint("environment-installation", encoded) - if err != nil { - return "", 0, err - } - return encoded + "." + signature, claim.ExpiresAt, nil -} - -func (s *Store) ValidateEnvironmentInstallation(ctx context.Context, token, version string) (sessions.InstallationAuthorization, error) { - var claim sessions.InstallationAuthorization - encoded, signature, ok := strings.Cut(token, ".") - if !ok || len(token) > 4096 { - return claim, sessions.ErrInstallationAuthorization - } - want, err := s.credentialCipher.Fingerprint("environment-installation", encoded) - if err != nil || !hmac.Equal([]byte(want), []byte(signature)) { - return claim, sessions.ErrInstallationAuthorization - } - payload, err := base64.RawURLEncoding.DecodeString(encoded) - if err != nil || json.Unmarshal(payload, &claim) != nil || claim.Version != version || claim.ExpiresAt <= time.Now().Unix() { - return sessions.InstallationAuthorization{}, sessions.ErrInstallationAuthorization - } - if err := s.selfHostedExecutorTarget(ctx, claim.Principal, claim.Environment); err != nil { - return sessions.InstallationAuthorization{}, sessions.ErrInstallationAuthorization - } - if err := activeProject(ctx, s.queries, claim.Principal); err != nil { - return sessions.InstallationAuthorization{}, sessions.ErrInstallationAuthorization - } - return claim, nil -} - -// ClaimEnvironmentInstallation stores only the digest of a secret generated and -// persisted by the installer before this request. A lost HTTP response is safe -// to retry; another machine or a revoked/rotated key cannot claim it again. -func (s *Store) ClaimEnvironmentInstallation(ctx context.Context, token, version, secret string) error { - claim, err := s.ValidateEnvironmentInstallation(ctx, token, version) - if err != nil { - return err - } - decoded, err := base64.RawURLEncoding.DecodeString(secret) - if err != nil || len(decoded) != 32 || base64.RawURLEncoding.EncodeToString(decoded) != secret { - return sessions.ErrInvalidInput - } - principal := claim.Principal - tenant, id, err := executorCredentialIdentity(principal, claim.Environment) - if err != nil { - return err - } - hash := sha256.Sum256([]byte(secret)) - digest := hex.EncodeToString(hash[:]) - return s.withExecutorCredentialTarget(ctx, principal, id, func(ctx context.Context, q *sqlc.Queries) error { - if err := activeProject(ctx, q, principal); err != nil { - return err - } - keys, err := q.ListEnvironmentExecutorCredentials(ctx, sqlc.ListEnvironmentExecutorCredentialsParams{TenantID: tenant, EnvironmentID: id, SubjectKind: pgtype.Text{String: principal.SubjectKind, Valid: true}, SubjectID: pgtype.Text{String: principal.SubjectID, Valid: true}}) - if err != nil { - return err - } - if len(keys) > 0 { - // Existing operator-issued keys must not be replaced by onboarding. - if len(keys) != 1 || keys[0].KeyID != id || keys[0].RevokedAt.Valid { - return sessions.ErrExecutorCredentialExists - } - _, err := q.AuthenticateEnvironmentExecutor(ctx, sqlc.AuthenticateEnvironmentExecutorParams{EnvironmentID: id, TokenSha256: digest}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrExecutorCredentialExists - } - return err - } - _, err = q.IssueExecutorCredential(ctx, sqlc.IssueExecutorCredentialParams{KeyID: id, TenantID: tenant, SubjectKind: pgtype.Text{String: principal.SubjectKind, Valid: true}, SubjectID: pgtype.Text{String: principal.SubjectID, Valid: true}, OrganizationID: principal.OrganizationID, ProjectID: principal.ProjectID, EnvironmentID: id, TokenSha256: digest}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrExecutorCredentialExists - } - return err - }) -} diff --git a/services/core/internal/store/environment_installation_test.go b/services/core/internal/store/environment_installation_test.go index a264db7ac..20c4726e4 100644 --- a/services/core/internal/store/environment_installation_test.go +++ b/services/core/internal/store/environment_installation_test.go @@ -22,6 +22,7 @@ func TestEnvironmentInstallationClaimLifetimeAndRetries(t *testing.T) { t.Fatal(err) } s := NewWithCredentialCipher(pool, cipher) + installations := sessionService(t, s) ctx := t.Context() p := createTestProject(t, pool).Principal input := environmentInput(uuid.NewString(), "self_hosted", "/workspace") @@ -34,12 +35,12 @@ func TestEnvironmentInstallationClaimLifetimeAndRetries(t *testing.T) { if err != nil { t.Fatal(err) } - token, expires, err := s.AuthorizeEnvironmentInstallation(ctx, p, environment.ID, "build") + token, expires, err := installations.AuthorizeEnvironmentInstallation(ctx, p, environment.ID, "build") if err != nil || expires <= time.Now().Unix() || expires > time.Now().Add(31*time.Minute).Unix() { t.Fatal("authorization", err) } for _, pair := range [][2]string{{token + "x", "build"}, {token, "other-build"}, {"", "build"}} { - if _, err := s.ValidateEnvironmentInstallation(ctx, pair[0], pair[1]); !errors.Is(err, sessions.ErrInstallationAuthorization) { + if _, err := installations.ValidateEnvironmentInstallation(ctx, pair[0], pair[1]); !errors.Is(err, sessions.ErrInstallationAuthorization) { t.Fatal("accepted invalid authorization", err) } } @@ -51,7 +52,7 @@ func TestEnvironmentInstallationClaimLifetimeAndRetries(t *testing.T) { raw, _ = json.Marshal(expired) payload = base64.RawURLEncoding.EncodeToString(raw) signature, _ := cipher.Fingerprint("environment-installation", payload) - if _, err := s.ValidateEnvironmentInstallation(ctx, payload+"."+signature, "build"); !errors.Is(err, sessions.ErrInstallationAuthorization) { + if _, err := installations.ValidateEnvironmentInstallation(ctx, payload+"."+signature, "build"); !errors.Is(err, sessions.ErrInstallationAuthorization) { t.Fatal("accepted expired grant", err) } one, _, _ := newExecutorSecret() @@ -61,7 +62,10 @@ func TestEnvironmentInstallationClaimLifetimeAndRetries(t *testing.T) { var wg sync.WaitGroup for i := range secrets { wg.Add(1) - go func() { defer wg.Done(); results[i] = s.ClaimEnvironmentInstallation(ctx, token, "build", secrets[i]) }() + go func() { + defer wg.Done() + results[i] = installations.ClaimEnvironmentInstallation(ctx, token, "build", secrets[i]) + }() } wg.Wait() winner := -1 @@ -78,22 +82,22 @@ func TestEnvironmentInstallationClaimLifetimeAndRetries(t *testing.T) { if winner < 0 { t.Fatal("no claim succeeded") } - if err := s.ClaimEnvironmentInstallation(ctx, token, "build", secrets[winner]); err != nil { + if err := installations.ClaimEnvironmentInstallation(ctx, token, "build", secrets[winner]); err != nil { t.Fatal("lost-response retry", err) } 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 { + if err := installations.RevokeExecutorCredential(ctx, p, environment.ID); err != nil { t.Fatal(err) } - if err := s.ClaimEnvironmentInstallation(ctx, token, "build", secrets[winner]); !errors.Is(err, sessions.ErrExecutorCredentialExists) { + if err := installations.ClaimEnvironmentInstallation(ctx, token, "build", secrets[winner]); !errors.Is(err, sessions.ErrExecutorCredentialExists) { t.Fatal("revoked key resurrected", err) } if err := s.DeleteSession(ctx, p.TenantID, session.ID); err != nil { t.Fatal(err) } - if _, err := s.ValidateEnvironmentInstallation(ctx, token, "build"); !errors.Is(err, sessions.ErrInstallationAuthorization) { + if _, err := installations.ValidateEnvironmentInstallation(ctx, token, "build"); !errors.Is(err, sessions.ErrInstallationAuthorization) { t.Fatal("deleted Session grant accepted", err) } } diff --git a/services/core/internal/store/environment_retrieve_public_test.go b/services/core/internal/store/environment_retrieve_public_test.go index 658d013e6..4fb1b5214 100644 --- a/services/core/internal/store/environment_retrieve_public_test.go +++ b/services/core/internal/store/environment_retrieve_public_test.go @@ -34,7 +34,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { if err := s.EnsureProjectScopes(t.Context(), []identity.ProjectScope{{TenantID: tenant, OrganizationID: principal.OrganizationID, ProjectID: tenant}, {TenantID: foreignTenant, OrganizationID: principal.OrganizationID, ProjectID: foreignTenant}}); err != nil { t.Fatal(err) } - executor, err := s.IssueExecutorCredential(t.Context(), principal, uuid.NewString(), "") + executor, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), "") if err != nil { t.Fatal(err) } @@ -43,7 +43,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { if !revoked { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if err := s.RevokeExecutorCredential(ctx, principal, executor.KeyID); err != nil { + if err := fixtureSessionService(t, db).RevokeExecutorCredential(ctx, principal, executor.KeyID); err != nil { t.Error("owned executor credential cleanup failed", err) } } @@ -80,7 +80,7 @@ func TestEnvironmentRetrievalOfficialClient(t *testing.T) { if err != nil { t.Fatal(err) } - if err := s.RevokeExecutorCredential(t.Context(), principal, executor.KeyID); err != nil { + if err := fixtureSessionService(t, db).RevokeExecutorCredential(t.Context(), principal, executor.KeyID); err != nil { t.Fatal(err) } revoked = true diff --git a/services/core/internal/store/environment_runtime_fixture_test.go b/services/core/internal/store/environment_runtime_fixture_test.go index 2671ce100..1366e4d2f 100644 --- a/services/core/internal/store/environment_runtime_fixture_test.go +++ b/services/core/internal/store/environment_runtime_fixture_test.go @@ -21,7 +21,7 @@ func enrollFixtureSession(t *testing.T, s *store.Store, db fixtureDB, tenant str t.Fatal(err) } principal := store.FixtureExecutorPrincipal(t, s, tenant) - key, err := s.IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + key, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_write_audit_test.go b/services/core/internal/store/environment_write_audit_test.go index 4603cda5b..eeb8afbf6 100644 --- a/services/core/internal/store/environment_write_audit_test.go +++ b/services/core/internal/store/environment_write_audit_test.go @@ -23,12 +23,12 @@ func TestEnvironmentUploadWriteAuditSurvivesRequestAndLeaseContext(t *testing.T) cancel() sessionAuditCount(t, f.s, f.tenant, 0) // Neither the settlement caller nor the newly acquired execution lease owns the request context. - awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, f.writer.pool) - if err := f.writer.lease.Close(context.Background()); err != nil { + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, f.s.pool) + if err := f.lease.Close(context.Background()); err != nil { t.Fatal(err) } awaitRelease() - next := executionWriter(t, f.s) + next := sessionExecution(t, executionWriter(t, f.s).lease) var group sync.WaitGroup for range 4 { group.Go(func() { @@ -62,7 +62,7 @@ func TestEnvironmentUploadWriteAuditRollbackAndRejection(t *testing.T) { if rejected && err != nil || !rejected && err == nil { t.Fatal("unexpected settlement outcome", err) } - got, err := f.s.GetEnvironmentFileWrite(t.Context(), f.tenant, f.env.ID, f.key.ID) + got, err := FixtureFileWrite(t.Context(), f.s.pool, f.tenant, f.env.ID, f.key.ID) want := "pending" if rejected { want = "rejected" diff --git a/services/core/internal/store/executor_credential_target.go b/services/core/internal/store/executor_credential_target.go deleted file mode 100644 index 3a3796866..000000000 --- a/services/core/internal/store/executor_credential_target.go +++ /dev/null @@ -1,70 +0,0 @@ -package store - -import ( - "context" - "crypto/rand" - "crypto/sha256" - "encoding/base64" - "encoding/hex" - "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/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" -) - -func executorCredentialIdentity(principal identity.Principal, keyID string) (pgtype.UUID, pgtype.UUID, error) { - if err := principal.Validate(); err != nil { - return pgtype.UUID{}, pgtype.UUID{}, fmt.Errorf("%w: %v", sessions.ErrInvalidInput, err) - } - tenant, err := parseID(principal.TenantID) - if err != nil { - return pgtype.UUID{}, pgtype.UUID{}, err - } - id, err := parseID(keyID) - return tenant, id, err -} - -func newExecutorSecret() (string, string, error) { - secret := make([]byte, 32) - if _, err := rand.Read(secret); err != nil { - return "", "", err - } - token := base64.RawURLEncoding.EncodeToString(secret) - hash := sha256.Sum256([]byte(token)) - return token, hex.EncodeToString(hash[:]), nil -} - -func issuedExecutorCredential(id, environment pgtype.UUID, token string) sessions.IssuedExecutorCredential { - result := sessions.IssuedExecutorCredential{KeyID: uuid.UUID(id.Bytes).String(), Token: token} - if environment.Valid { - result.EnvironmentID = uuid.UUID(environment.Bytes).String() - } - return result -} - -func (s *Store) withExecutorCredentialTarget(ctx context.Context, principal identity.Principal, environment pgtype.UUID, apply func(context.Context, *sqlc.Queries) error) error { - 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 := sessionpg.LoadEnvironment(ctx, s.queries, principal.TenantID, uuid.UUID(environment.Bytes).String()) - if err != nil { - return err - } - tenant, _ := parseID(principal.TenantID) - // The Session lock orders exact-target issuance and rotation against deletion. - return s.withPublicSession(ctx, principal.TenantID, owned.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - row, err := q.GetSession(ctx, sqlc.GetSessionParams{TenantID: tenant, ID: session}) - if err != nil { - return err - } - if !row.CreatorKind.Valid || !row.CreatorID.Valid || row.CreatorKind.String != principal.SubjectKind || row.CreatorID.String != principal.SubjectID { - return sessions.ErrNotFound - } - return apply(ctx, q) - }) -} diff --git a/services/core/internal/store/executor_principals_migration_test.go b/services/core/internal/store/executor_principals_migration_test.go index b2d5bc876..258634aa3 100644 --- a/services/core/internal/store/executor_principals_migration_test.go +++ b/services/core/internal/store/executor_principals_migration_test.go @@ -105,26 +105,27 @@ func TestExecutorPrincipalMigrationRetiresUnknownAuthority(t *testing.T) { t.Cleanup(migrated.Close) s := New(migrated) p := FixtureExecutorPrincipal(t, s, tenant) + credentials := sessionService(t, s) 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) { + if _, err := credentials.IssueExecutorCredential(ctx, p, environment, ""); !errors.Is(err, sessions.ErrExecutorCredentialExists) { t.Fatal("legacy key ID claimed", err) } - if _, err := s.RotateExecutorCredential(ctx, p, environment); !errors.Is(err, sessions.ErrNotFound) { + if _, err := credentials.RotateExecutorCredential(ctx, p, environment); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("legacy principal claimed", err) } - if err := s.RevokeExecutorCredential(ctx, p, environment); !errors.Is(err, sessions.ErrNotFound) { + if err := credentials.RevokeExecutorCredential(ctx, p, environment); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("legacy principal manufactured", err) } - key, err := s.IssueExecutorCredential(ctx, p, uuid.NewString(), "") + key, err := credentials.IssueExecutorCredential(ctx, p, uuid.NewString(), "") if err != nil { t.Fatal(err) } if _, err := provider.DownTo(ctx, 25); err == nil || !strings.Contains(err.Error(), "Cannot remove durable executor principal identities") { t.Fatal("downgrade lost principal keys", err) } - if _, err := s.RotateExecutorCredential(ctx, p, key.KeyID); err != nil { + if _, err := credentials.RotateExecutorCredential(ctx, p, key.KeyID); err != nil { t.Fatal("failed downgrade damaged identity", err) } } diff --git a/services/core/internal/store/executor_principals_test.go b/services/core/internal/store/executor_principals_test.go index 3b9ecc4b8..3754e390c 100644 --- a/services/core/internal/store/executor_principals_test.go +++ b/services/core/internal/store/executor_principals_test.go @@ -13,8 +13,9 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.T) { s, pool := testStore(t) ctx := t.Context() p := FixtureExecutorPrincipal(t, s, uuid.NewString()) + credentials := sessionService(t, s) keyID := uuid.NewString() - issued, err := s.IssueExecutorCredential(ctx, p, keyID, "") + issued, err := credentials.IssueExecutorCredential(ctx, p, keyID, "") if err != nil || issued.KeyID != keyID || issued.EnvironmentID != "" || len(issued.Token) != 43 { t.Fatal("pre-Session principal issuance failed", err) } @@ -55,13 +56,13 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.T) { for _, different := range []identity.Principal{otherKind, otherID, foreign} { _, target := create(different) check(s, target.ID, issued, false) - if _, err := s.IssueExecutorCredential(ctx, p, uuid.NewString(), target.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := credentials.IssueExecutorCredential(ctx, p, uuid.NewString(), target.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("restricted key accepted another creator/project", err) } - if _, err := s.RotateExecutorCredential(ctx, different, keyID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := credentials.RotateExecutorCredential(ctx, different, keyID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign rotation", err) } - if err := s.RevokeExecutorCredential(ctx, different, keyID); !errors.Is(err, sessions.ErrNotFound) { + if err := credentials.RevokeExecutorCredential(ctx, different, keyID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign revocation", err) } } @@ -69,17 +70,17 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.T) { {ProjectScope: identity.ProjectScope{TenantID: p.TenantID, OrganizationID: "other-org", ProjectID: p.ProjectID}, SubjectKind: p.SubjectKind, SubjectID: p.SubjectID}, {ProjectScope: identity.ProjectScope{TenantID: p.TenantID, OrganizationID: p.OrganizationID, ProjectID: "other-project"}, SubjectKind: p.SubjectKind, SubjectID: p.SubjectID}, } { - if _, err := s.IssueExecutorCredential(ctx, different, uuid.NewString(), ""); !errors.Is(err, sessions.ErrNotFound) { + if _, err := credentials.IssueExecutorCredential(ctx, different, uuid.NewString(), ""); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unverified scope issuance", err) } - if _, err := s.RotateExecutorCredential(ctx, different, keyID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := credentials.RotateExecutorCredential(ctx, different, keyID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unverified scope rotation", err) } - if err := s.RevokeExecutorCredential(ctx, different, keyID); !errors.Is(err, sessions.ErrNotFound) { + if err := credentials.RevokeExecutorCredential(ctx, different, keyID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unverified scope revocation", err) } } - restricted, err := s.IssueExecutorCredential(ctx, p, uuid.NewString(), one.ID) + restricted, err := credentials.IssueExecutorCredential(ctx, p, uuid.NewString(), one.ID) if err != nil { t.Fatal(err) } @@ -88,6 +89,7 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.T) { // Reopening the database preserves both key identity and multi-Session authority. pool.Close() restarted, newPool := testStore(t) + credentials = sessionService(t, restarted) check(restarted, one.ID, issued, true) check(restarted, two.ID, issued, true) if err := restarted.DeleteSession(ctx, p.TenantID, first.ID); err != nil { @@ -95,10 +97,10 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.T) { } check(restarted, one.ID, issued, false) check(restarted, two.ID, issued, true) - if err := restarted.RevokeExecutorCredential(ctx, p, restricted.KeyID); err != nil { + if err := credentials.RevokeExecutorCredential(ctx, p, restricted.KeyID); err != nil { t.Fatal("revoke deleted restriction", err) } - rotated, err := restarted.RotateExecutorCredential(ctx, p, keyID) + rotated, err := credentials.RotateExecutorCredential(ctx, p, keyID) if err != nil || rotated.KeyID != keyID || rotated.EnvironmentID != "" || rotated.Token == issued.Token { t.Fatal("principal rotation", err) } @@ -110,14 +112,14 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.T) { FROM environment_executor_credentials WHERE key_id=$1`, keyID, p.TenantID, p.SubjectKind, p.SubjectID).Scan(&retained); err != nil || !retained { t.Fatal("rotation changed principal or creation identity", err) } - if err := restarted.RevokeExecutorCredential(ctx, p, keyID); err != nil { + if err := credentials.RevokeExecutorCredential(ctx, p, keyID); err != nil { t.Fatal(err) } check(restarted, two.ID, rotated, false) - if _, err := restarted.IssueExecutorCredential(ctx, p, keyID, ""); !errors.Is(err, sessions.ErrExecutorCredentialExists) { + if _, err := credentials.IssueExecutorCredential(ctx, p, keyID, ""); !errors.Is(err, sessions.ErrExecutorCredentialExists) { t.Fatal("issue restored revoked authority", err) } - restored, err := restarted.RotateExecutorCredential(ctx, p, keyID) + restored, err := credentials.RotateExecutorCredential(ctx, p, keyID) if err != nil { t.Fatal(err) } @@ -128,14 +130,15 @@ func TestExecutorPrincipalRequiresVerifiedScopeAndRecordedCreator(t *testing.T) s, pool := testStore(t) ctx := t.Context() p := identity.Principal{ProjectScope: identity.ProjectScope{TenantID: uuid.NewString(), OrganizationID: "org", ProjectID: uuid.NewString()}, SubjectKind: "user", SubjectID: "owner"} - if _, err := s.IssueExecutorCredential(ctx, p, uuid.NewString(), ""); !errors.Is(err, sessions.ErrNotFound) { + credentials := sessionService(t, s) + if _, err := credentials.IssueExecutorCredential(ctx, p, uuid.NewString(), ""); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("issuer manufactured a project mapping", err) } if err := s.EnsureProjectScopes(ctx, []identity.ProjectScope{p.ProjectScope}); err != nil { t.Fatal(err) } for _, invalid := range []identity.Principal{{}, {ProjectScope: p.ProjectScope}, {ProjectScope: p.ProjectScope, SubjectKind: "workspace", SubjectID: "owner"}} { - if _, err := s.IssueExecutorCredential(ctx, invalid, uuid.NewString(), ""); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := credentials.IssueExecutorCredential(ctx, invalid, uuid.NewString(), ""); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal("invalid principal", err) } } @@ -149,7 +152,7 @@ func TestExecutorPrincipalRequiresVerifiedScopeAndRecordedCreator(t *testing.T) if err != nil { t.Fatal(err) } - key, err := s.IssueExecutorCredential(ctx, p, uuid.NewString(), "") + key, err := credentials.IssueExecutorCredential(ctx, p, uuid.NewString(), "") if err != nil { t.Fatal(err) } @@ -162,7 +165,7 @@ func TestExecutorPrincipalRequiresVerifiedScopeAndRecordedCreator(t *testing.T) 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) { + if _, err := credentials.IssueExecutorCredential(ctx, p, uuid.NewString(), target.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unknown creator claimed", err) } } diff --git a/services/core/internal/store/export_test.go b/services/core/internal/store/export_test.go index 4dd9d8b7b..fc18300a2 100644 --- a/services/core/internal/store/export_test.go +++ b/services/core/internal/store/export_test.go @@ -12,6 +12,7 @@ import ( "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/identity" + "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" @@ -106,6 +107,24 @@ func FixtureFunctionCall(ctx context.Context, pool *pgxpool.Pool, tenantID, sess return sessions.FunctionCall{CallID: row.CallID, ExecutorCallID: row.ExecutorCallID, Name: row.Name, Arguments: row.Arguments, Result: row.Result, Applied: row.Applied}, nil } +// FixtureFileWrite reads a file write to the tenant's Environment from the +// test's database for assertions, including one whose Session was publicly +// deleted. An unknown write is sessions.ErrNotFound. +func FixtureFileWrite(ctx context.Context, pool *pgxpool.Pool, tenantID, environmentID, writeID string) (sessions.EnvironmentFileWrite, error) { + row, err := sqlc.New(pool).GetEnvironmentFileWrite(ctx, sqlc.GetEnvironmentFileWriteParams{TenantID: pgunit.PathID(tenantID), EnvironmentID: pgunit.PathID(environmentID), ID: pgunit.PathID(writeID)}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.EnvironmentFileWrite{}, sessions.ErrNotFound + } + if err != nil { + return sessions.EnvironmentFileWrite{}, err + } + write := row.EnvironmentFileWrite + return sessions.EnvironmentFileWrite{ + Identity: sessions.FileWriteIdentity{ID: uuid.UUID(write.ID.Bytes).String(), DeviceID: uuid.UUID(write.DeviceID.Bytes).String(), RequestSHA256: write.RequestSha256}, + EnvironmentID: uuid.UUID(write.EnvironmentID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), State: write.State, CreatedAt: write.CreatedAt.Time, + }, nil +} + // SubmitFixtureFunctionResult submits result for the Turn's call as a public // tool_result input under a fresh request key. func SubmitFixtureFunctionResult(ctx context.Context, s *Store, tenantID, sessionID, turnID, callID string, result json.RawMessage) error { diff --git a/services/core/internal/store/local_environment_file_write_test.go b/services/core/internal/store/local_environment_file_write_test.go index d08028fc7..1f255c6e7 100644 --- a/services/core/internal/store/local_environment_file_write_test.go +++ b/services/core/internal/store/local_environment_file_write_test.go @@ -9,6 +9,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -52,7 +53,7 @@ func TestLocalEnvironmentFileWriteOwnsMutationBeforeDispatch(t *testing.T) { if begin.DecodePayload(&request) != nil || request.Step != "begin" || request.EnvironmentID != environment.ID || request.SessionID != h.session.ID || request.Path != "input" || request.SizeBytes != 3 { t.Fatal("upload identity changed") } - intent, err := h.s.GetEnvironmentFileWrite(t.Context(), h.tenant, environment.ID, begin.ID) + intent, err := store.FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, begin.ID) if err != nil || intent.State != "pending" || intent.Identity.DeviceID != h.device.ID { t.Fatal("dispatch preceded durable ownership", intent, err) } @@ -75,7 +76,7 @@ func TestLocalEnvironmentFileWriteOwnsMutationBeforeDispatch(t *testing.T) { } h.write(begin.ID, proto.TypeWorkspaceWriteResult, proto.WorkspaceWriteResultPayload{Outcome: "completed", SizeBytes: 3}) awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "detached durable commit", func() bool { - got, e := h.s.GetEnvironmentFileWrite(t.Context(), h.tenant, environment.ID, begin.ID) + got, e := store.FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, begin.ID) return e == nil && got.State == "committed" }) session, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) @@ -94,7 +95,7 @@ func TestLocalEnvironmentFileWriteLostReceiptRemainsPending(t *testing.T) { if result := awaitLocalWrite(t, done); !errors.Is(result.err, execution.ErrExecutionUnavailable) { t.Fatal(result.err) } - intent, err := h.s.GetEnvironmentFileWrite(t.Context(), h.tenant, environment.ID, begin.ID) + intent, err := store.FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, begin.ID) if err != nil || intent.State != "pending" { t.Fatal("disconnect guessed rejection", intent, err) } @@ -112,7 +113,7 @@ func TestLocalEnvironmentFileWriteKnownRejectionReleasesMutation(t *testing.T) { if result := awaitLocalWrite(t, done); !errors.Is(result.err, execution.ErrExecutionUnavailable) { t.Fatal(result.err) } - intent, err := h.s.GetEnvironmentFileWrite(t.Context(), h.tenant, environment.ID, begin.ID) + intent, err := store.FixtureFileWrite(t.Context(), h.db.pool, h.tenant, environment.ID, begin.ID) if err != nil || intent.State != "rejected" { t.Fatal("rejection did not settle", intent, err) } diff --git a/services/core/internal/store/public_handler_fixture_test.go b/services/core/internal/store/public_handler_fixture_test.go index 2dc3288c8..923187664 100644 --- a/services/core/internal/store/public_handler_fixture_test.go +++ b/services/core/internal/store/public_handler_fixture_test.go @@ -93,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, EnvironmentsReader: sessionStore, Admin: s, AdminAudit: audit, WriteAudit: audit, + SessionAdmin: s, Environments: sessionService, 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_enrollment_connection_test.go b/services/core/internal/store/runtime_enrollment_connection_test.go index 99d805920..aea56cc8c 100644 --- a/services/core/internal/store/runtime_enrollment_connection_test.go +++ b/services/core/internal/store/runtime_enrollment_connection_test.go @@ -36,7 +36,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { if err != nil { t.Fatal(err) } - key, err := s.IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + key, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } @@ -70,13 +70,13 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { } assertConnection(environment.ID, key.Token, "disconnected", 200) assertConnection(uuid.NewString(), key.Token, "", 401) - otherKey, err := s.IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) + otherKey, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), principal, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } assertConnection(environment.ID, otherKey.Token, "", 409) foreign := store.FixtureExecutorPrincipal(t, s, uuid.NewString()) - foreignKey, err := s.IssueExecutorCredential(t.Context(), foreign, uuid.NewString(), "") + foreignKey, err := fixtureSessionService(t, db).IssueExecutorCredential(t.Context(), foreign, uuid.NewString(), "") if err != nil { t.Fatal(err) } @@ -132,7 +132,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { first := connect(key.Token) await("connected") assertConnection(environment.ID, key.Token, "connected", 200) - rotated, err := s.RotateExecutorCredential(t.Context(), principal, key.KeyID) + rotated, err := fixtureSessionService(t, db).RotateExecutorCredential(t.Context(), principal, key.KeyID) if err != nil { t.Fatal(err) } @@ -155,7 +155,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { // live, authorized daemon. No compute allocation or native execution is made. stop = start() await("connected") - if err = s.RevokeExecutorCredential(t.Context(), principal, key.KeyID); err != nil { + if err = fixtureSessionService(t, db).RevokeExecutorCredential(t.Context(), principal, key.KeyID); err != nil { t.Fatal(err) } assertConnection(environment.ID, rotated.Token, "", 401) diff --git a/services/core/internal/store/runtime_enrollment_test.go b/services/core/internal/store/runtime_enrollment_test.go index 8384f2f1a..2a4a001ba 100644 --- a/services/core/internal/store/runtime_enrollment_test.go +++ b/services/core/internal/store/runtime_enrollment_test.go @@ -24,7 +24,7 @@ func runtimeEnrollmentFixture(t *testing.T, s *Store, p identity.Principal) (ses if err != nil { t.Fatal(err) } - key, err := s.IssueExecutorCredential(t.Context(), p, uuid.NewString(), environment.ID) + key, err := sessionService(t, s).IssueExecutorCredential(t.Context(), p, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } @@ -60,14 +60,14 @@ func TestRuntimeEnrollmentAuthorityAndRotation(t *testing.T) { if _, err := auth.AuthenticateBearer(ctx, bound.DeviceID, key.Token); err != nil { t.Fatal(err) } - otherKey, err := s.IssueExecutorCredential(ctx, p, uuid.NewString(), environment.ID) + otherKey, err := sessionService(t, s).IssueExecutorCredential(ctx, p, uuid.NewString(), environment.ID) if err != nil { t.Fatal(err) } 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) + rotated, err := sessionService(t, s).RotateExecutorCredential(ctx, p, key.KeyID) if err != nil { t.Fatal(err) } @@ -89,7 +89,7 @@ func TestRuntimeEnrollmentAuthorityAndRotation(t *testing.T) { 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 { + if err := sessionService(t, s).RevokeExecutorCredential(ctx, p, key.KeyID); err != nil { t.Fatal(err) } if _, ok, err := sessionAdapter(s).GetDeviceCredential(ctx, bound.DeviceID); err != nil || ok { diff --git a/services/core/internal/store/runtime_file_admission_test.go b/services/core/internal/store/runtime_file_admission_test.go index c647eec26..d50b8092f 100644 --- a/services/core/internal/store/runtime_file_admission_test.go +++ b/services/core/internal/store/runtime_file_admission_test.go @@ -19,25 +19,26 @@ func runtimeFileWriteKey(owner RuntimeAllocation) sessions.FileWriteIdentity { func TestRuntimeFileWriteRequiresRunningComputeBeforeNewIntent(t *testing.T) { for _, phase := range []string{"disabled", "running", "quiescing", "suspending", "suspended", "restoring", "waking"} { t.Run(phase, func(t *testing.T) { - s, w, pool, owner := runtimeSuspensionFixture(t) + _, w, pool, owner := runtimeSuspensionFixture(t) + writes := sessionExecution(t, w.lease) runtimeSuspensionSQL(t, pool, `UPDATE runtime_allocations SET compute_phase=$2,compute_retained_until=clock_timestamp()+interval '1 hour' WHERE id=$1`, owner.ID, phase) key := runtimeFileWriteKey(owner) - write, err := w.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key) + write, err := writes.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key) blocked := phase != "disabled" && phase != "running" if blocked { if !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("suspended compute admitted a new write", write, err) } - if _, err := s.GetEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := FixtureFileWrite(t.Context(), pool, owner.TenantID, owner.EnvironmentID, key.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("rejected admission retained a blocking intent", err) } runtimeSuspensionSQL(t, pool, `UPDATE runtime_allocations SET compute_phase='running',compute_retained_until=NULL WHERE id=$1`, owner.ID) - write, err = w.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key) + write, err = writes.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key) } if err != nil || write.Replayed || write.State != "pending" { t.Fatal("running compute could not admit original request", write, err) } - if _, err := w.SettleEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key, "committed"); err != nil { + if _, err := writes.SettleEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key, "committed"); err != nil { t.Fatal(err) } }) @@ -46,8 +47,9 @@ func TestRuntimeFileWriteRequiresRunningComputeBeforeNewIntent(t *testing.T) { func TestRuntimeFileWritePhaseFencePreservesExistingReceipts(t *testing.T) { _, w, pool, owner := runtimeSuspensionFixture(t) + writes := sessionExecution(t, w.lease) key := runtimeFileWriteKey(owner) - first, err := w.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key) + first, err := writes.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key) if err != nil { t.Fatal(err) } @@ -55,20 +57,20 @@ func TestRuntimeFileWritePhaseFencePreservesExistingReceipts(t *testing.T) { // it never grants permission to send the unknown write again. runtimeSuspensionSQL(t, pool, `UPDATE runtime_allocations SET compute_phase='waking',compute_retained_until=clock_timestamp()+interval '1 hour' WHERE id=$1`, owner.ID) for _, state := range []string{"pending", "committed"} { - got, err := w.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key) + got, err := writes.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key) if err != nil || !got.Replayed || got.State != state || got.Identity != key || !got.CreatedAt.Equal(first.CreatedAt) { t.Fatal("phase fence changed the existing receipt", got, err) } changed := key changed.RequestSHA256 = strings.Repeat("b", 64) - if _, err := w.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { + if _, err := writes.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal("phase fence masked changed retry identity", err) } - if _, err := w.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, runtimeFileWriteKey(owner)); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := writes.ReserveEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, runtimeFileWriteKey(owner)); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("receipt authorized a successor while waking", err) } if state == "pending" { - if _, err := w.SettleEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key, "committed"); err != nil { + if _, err := writes.SettleEnvironmentFileWrite(t.Context(), owner.TenantID, owner.EnvironmentID, key, "committed"); err != nil { t.Fatal("phase fence prevented exact receipt settlement", err) } } @@ -80,6 +82,7 @@ func TestRuntimeFileWriteAndQuiesceSerializeBothOrders(t *testing.T) { t.Run(first, func(t *testing.T) { s, w, pool, owner := runtimeSuspensionFixture(t) runtimeSuspensionCompleted(t, pool, owner) + writes := sessionExecution(t, w.lease) key := runtimeFileWriteKey(owner) until := time.Now().Add(time.Hour) ctx, tx, blocker := runtimeSuspensionLockedSession(t, pool, owner.SessionID) @@ -87,7 +90,7 @@ func TestRuntimeFileWriteAndQuiesceSerializeBothOrders(t *testing.T) { go func() { var err error if first == "quiesce" { - _, err = w.ReserveEnvironmentFileWrite(ctx, owner.TenantID, owner.EnvironmentID, key) + _, err = writes.ReserveEnvironmentFileWrite(ctx, owner.TenantID, owner.EnvironmentID, key) } else { _, err = w.SetRuntimeCompute(ctx, owner, "quiescing", json.RawMessage(`{}`), &until, time.Nanosecond) } @@ -115,7 +118,7 @@ func TestRuntimeFileWriteAndQuiesceSerializeBothOrders(t *testing.T) { if err != nil { t.Fatal(err) } - write, err := s.GetEnvironmentFileWrite(ctx, owner.TenantID, owner.EnvironmentID, key.ID) + write, err := FixtureFileWrite(ctx, pool, owner.TenantID, owner.EnvironmentID, key.ID) if first == "quiesce" { if allocation.ComputePhase != "quiescing" || !errors.Is(err, sessions.ErrNotFound) { t.Fatal("late write survived quiesce", allocation.ComputePhase, write, err) diff --git a/services/core/internal/store/sessions_fixture_test.go b/services/core/internal/store/sessions_fixture_test.go index bd878fbc4..ff2aaba29 100644 --- a/services/core/internal/store/sessions_fixture_test.go +++ b/services/core/internal/store/sessions_fixture_test.go @@ -30,7 +30,7 @@ func fixtureSessionStore(db fixtureDB) *sessionpg.Store { } // fixtureSessionService is the Session service on fixtureSessionStore(db). It -// runs the device, heartbeat and enrollment use cases. +// runs the device, heartbeat, enrollment and executor credential use cases. func fixtureSessionService(t testing.TB, db fixtureDB) *sessions.Service { t.Helper() service, err := sessions.NewService(fixtureSessionStore(db)) diff --git a/services/core/internal/store/sessions_test.go b/services/core/internal/store/sessions_test.go index 8227d2f89..4cb335a57 100644 --- a/services/core/internal/store/sessions_test.go +++ b/services/core/internal/store/sessions_test.go @@ -27,7 +27,8 @@ func testStore(t *testing.T) (*Store, *pgxpool.Pool) { 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. +// device, heartbeat, enrollment, executor credential and installation use +// cases. func sessionService(t testing.TB, s *Store) *sessions.Service { t.Helper() service, err := sessions.NewService(sessionAdapter(s))