From f29e406e78f247659bdcb369cf6b7549e53a5b8d Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 00:10:38 +0800 Subject: [PATCH 1/6] Move Core metrics, Runtime history and Runtime targets out of store (#430) --- services/core/IMPLEMENTATION.md | 6 +- services/core/cmd/server/core_metrics.go | 18 +-- services/core/cmd/server/http_routes_test.go | 2 +- services/core/cmd/server/main.go | 9 +- services/core/cmd/server/runtime_history.go | 6 +- .../core/cmd/server/runtime_history_test.go | 6 +- services/core/internal/api/admin_resources.go | 4 +- services/core/internal/api/admin_runtime.go | 11 +- .../core/internal/api/admin_runtime_test.go | 24 ++-- services/core/internal/api/dependencies.go | 3 +- .../core/internal/api/dependencies_test.go | 5 +- services/core/internal/api/fakes_test.go | 12 +- services/core/internal/coremetrics/service.go | 3 +- services/core/internal/coremetrics/types.go | 15 +++ .../archive_cancellation_cleanup_test.go | 2 +- .../postgres/coremetricspg/store.go | 107 +++++++++++++++++ .../postgres/coremetricspg/store_test.go} | 45 +++---- .../persistence/postgres/sessionpg/admin.go | 55 +++++++++ .../postgres/sessionpg/admin_test.go | 63 ++++++++++ .../persistence/postgres/sessionpg/devices.go | 23 ++++ .../runtimehistory/postgresreader/reader.go | 15 ++- .../postgresreader/reader_test.go | 2 +- .../postgresreader/samples.go} | 39 +++--- .../postgresreader/samples_test.go} | 80 ++++++++----- .../providers/configuration_flow_test.go | 2 +- services/core/internal/sessions/admin.go | 23 ++++ services/core/internal/sessions/devices.go | 4 + services/core/internal/sessions/reader.go | 6 +- services/core/internal/store/admin_summary.go | 50 -------- .../store/archive_cancellation_test.go | 6 +- services/core/internal/store/core_metrics.go | 111 ------------------ services/core/internal/store/devices_test.go | 2 +- services/core/internal/store/dispatch_test.go | 2 +- .../internal/store/node_host_history_test.go | 23 ---- .../store/public_handler_fixture_test.go | 2 +- .../internal/store/runtime_cancellation.go | 35 ------ .../internal/store/runtime_connection_test.go | 2 +- .../runtime_enrollment_connection_test.go | 2 +- 38 files changed, 473 insertions(+), 352 deletions(-) create mode 100644 services/core/internal/persistence/postgres/coremetricspg/store.go rename services/core/internal/{store/core_metrics_test.go => persistence/postgres/coremetricspg/store_test.go} (79%) create mode 100644 services/core/internal/persistence/postgres/sessionpg/admin.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/admin_test.go rename services/core/internal/{store/runtime_history.go => runtimehistory/postgresreader/samples.go} (72%) rename services/core/internal/{store/runtime_history_acceptance_test.go => runtimehistory/postgresreader/samples_test.go} (76%) create mode 100644 services/core/internal/sessions/admin.go delete mode 100644 services/core/internal/store/core_metrics.go delete mode 100644 services/core/internal/store/node_host_history_test.go delete mode 100644 services/core/internal/store/runtime_cancellation.go diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 9231be3fa..29e210fc0 100644 --- a/services/core/IMPLEMENTATION.md +++ b/services/core/IMPLEMENTATION.md @@ -6,7 +6,7 @@ These are the code-level rules of `services/core` that no contract states. Contr `api.NewHandler` takes one `api.Dependencies` value, built only in `cmd/server`. Each application area is one field typed as an interface declared in `api` beside its handlers, listing exactly the methods they call. Every field is required and `NewHandler` rejects a missing one, except the optional groups whose comments say what nil means: `Execution` is nil without an execution Worker, `Sandboxes` is nil without a managed sandbox installation and requires `Execution`, and `Execution.NativeInstaller` is nil for a build without a source revision. Handlers never discover a capability by type assertion or fall back to another implementation. API tests use one strict fake per area, `fake`, which fails the test on any call the test did not set. -`internal/persistence/postgres/pgunit` owns the PostgreSQL mechanics that adapters share: pooled read-write and snapshot transactions; pool-bound queries, used only for a single-statement read that needs no transaction, such as the per-request key lookup; the execution lease (its dedicated connection and gate, the ownership check, the cancellation fence, close, and the execution deadline); identifier parsing (`ParseID`, `PathID`, `LookupCursor`); and detection of text PostgreSQL cannot store (`IsUnstorableText`). Persistence code runs every transaction through it. Outside `persistence` and `store`, only `cmd/server`, which acquires the lease and builds the adapters' pool, and test fixtures import it. `internal/persistence/postgres/pgtest` is test support: it opens the dedicated test database under the `oac_*_tests` guard, applies the migrations, and creates isolated databases for database-wide state such as the execution lease. Only test files import it. +`internal/persistence/postgres/pgunit` owns the PostgreSQL mechanics that adapters share: pooled read-write and snapshot transactions; pool-bound queries, used only for a single-statement read that needs no transaction, such as the per-request key lookup; the execution lease (its dedicated connection and gate, the ownership check, the cancellation fence, close, and the execution deadline); identifier parsing (`ParseID`, `PathID`, `LookupCursor`); and detection of text PostgreSQL cannot store (`IsUnstorableText`). Persistence code runs every transaction through it. Outside `persistence` and `store`, only `cmd/server`, which acquires the lease and builds the adapters' pool, `runtimehistory/postgresreader`, which stores Runtime history in Core's database, and test fixtures import it. `internal/persistence/postgres/pgtest` is test support: it opens the dedicated test database under the `oac_*_tests` guard, applies the migrations, and creates isolated databases for database-wide state such as the execution lease. Only test files import it. `internal/persistence/postgres/auditpg` is the one adapter that other adapters call directly. Audit rows are written inside the business transaction, so each adapter calls `RecordWriteAudit`, `RecordAdminMutation` or `RecordDeploymentMutation` with its own transaction's queries; the audit provenance travels in the context. `writeaudit` and `adminaudit` own the sources, their validation, `ErrInvalidSource` and the read models, and `auditpg.Store` serves the audit reads. @@ -34,7 +34,9 @@ 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` and its subpackage `deployment/placement` (`deploymentpg`, `placementpg`): the sandbox deployment and its nodes: provider configuration and the sealed credential, the specification and retained generations, setup, update and switch, node enrollment, identity and authentication, generation configuration, capacity, presence, status, host history, reset, and the counts of nodes and sandboxes bound to the public address; and hosted runtime allocations: reservation with the dedicated device, compute ownership and settlement, observation diagnostics, activity, compute phases, wake receipts and cleanup, with the reads that schedule, discover and authorize them. It interprets Sandbox Provider declarations through the `providers.Registry` it is given, whose lookups return typed errors. `cmd/server` builds that registry and calls it only to build a direct Provider and to discover a Provider's configuration; it takes each setup's mode and declared operations from the `Setup` that `deployment` returns. `deployment/placement` owns hosted admission and placement: `cmd/server` builds one `placement.Rules` from the registry and the public URL, both fixed while Core runs, and gives it to `deployment.Service` and to `store`; its pure decisions admit a hosted Session, choose its node, and admit an allocation's reserved node and a restore on it, and its errors keep one status, code and message through `writeStoreError` and `writeDeploymentError`. `placementpg` loads the facts those decisions read and applies them on the caller's transaction-bound queries; it has no Store or transaction runner. The only other reader is `store`: Session creation decides admission and placement with those rules and `placementpg` inside its own transaction. Deployment changes and allocation writes run through `deployment.ExecutionOperations` on `deploymentpg.NewExecution(lease, …)`, which the Worker receives as `execution.Owner.Deployment`. Each allocation write locks the owning Session, decides in `deployment`, applies in `deploymentpg` and prunes the Session's change journal in the same leased transaction; cleanup settles the Session through the `sessions` procedures on a `sessionpg.SessionTx` bound to that transaction. The Worker reads the deployment, prepares a selection's setup and records live-compute activity through the pooled `deployment.Service` in `execution.Dispatcher.Deployment`, and lists the Sessions a reset still has to archive, schedules node lifecycles and reads allocations through the pooled `deployment.Reader` in `execution.Dispatcher.DeploymentReader`; a pooled read carries no lease, and the leased write that follows it rechecks the allocation's owner. Node management and reads use the pooled `deploymentpg.Store`, which `cmd/server` also reads the owner epoch from and runs the host-history sampler on. Provider calls run outside transactions, and the final transaction rechecks the expected generation. The reset's Session archive stays in `store`. -- `sessions` (`sessionpg`): Session use cases and reads as they leave `store`. The pooled `sessions.Service` runs the use cases on `sessionpg.Store`, which implements `sessions.Storage`, and plain reads use `sessions.Reader`, which `sessionpg.Store` also implements, directly. So far these cover root Items and Subagents (their reads), Session Artifacts (their reads, deletion and the staging of a Turn's export), Environments (their reads, the initialization list and the frozen setup and initial files, which `sessionpg.Store` opens with the credential key), devices (their reads, creation, revocation, heartbeats and Runtime enrollment), executor credentials (authentication and a Project's credential state through `sessions.ExecutorCredentialReader`, and issuance, rotation and revocation; the Core-key Project operations record their audit entry in the same transaction) and native installation authorization and claims, whose tokens `sessionpg.Store` signs with the credential key. `cmd/server` builds one `sessionpg.Store` with the credential key and wires the Service and the Store into `execution.Dispatcher.Sessions` and `SessionsReader`, into the api fields, into Runtime enrollment and into the daemon gateway; the Worker stages Artifacts through the Service; `cmd/environment-key` builds one without the key for its credential commands. Function calls and their application receipts, the Turn execution journal, Environment initialization, connection observations and their reconciliation, Session device binding, and file-write reservation and settlement run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. +- `coremetrics` (`coremetricspg`): the Core metrics PostgreSQL holds: the root Turn queue counts, the root Turn history, read from one read-only snapshot, and the database size. `cmd/server`'s Core metrics source adds them to the process, pool, Worker and daemon registry measurements. +- `runtimehistory` (`runtimehistory/postgresreader`, outside this directory): Runtime history samples: the periodic export, scoped reads and retention, which also prunes node-host samples. +- `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 the archived cancellation receipt the daemon gateway reads), the Sessions whose Runtime the administrator observes across Projects, executor credentials (authentication and a Project's credential state through `sessions.ExecutorCredentialReader`, and issuance, rotation and revocation; the Core-key Project operations record their audit entry in the same transaction) and native installation authorization and claims, whose tokens `sessionpg.Store` signs with the credential key. `cmd/server` builds one `sessionpg.Store` with the credential key and wires the Service and the Store into `execution.Dispatcher.Sessions` and `SessionsReader`, into the api fields, into Runtime enrollment and into the daemon gateway; the Worker stages Artifacts through the Service; `cmd/environment-key` builds one without the key for its credential commands. Function calls and their application receipts, the Turn execution journal, Environment initialization, connection observations and their reconciliation, Session device binding, and file-write reservation and settlement run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. ## Request handling diff --git a/services/core/cmd/server/core_metrics.go b/services/core/cmd/server/core_metrics.go index 9c0147c39..84d4c59a4 100644 --- a/services/core/cmd/server/core_metrics.go +++ b/services/core/cmd/server/core_metrics.go @@ -6,8 +6,8 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/coremetrics" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/coremetricspg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/jackc/pgx/v5/pgxpool" ) @@ -16,7 +16,7 @@ var buildRevision string var processStartedAt = time.Now().UTC() type coreMetricsSource struct { - store *store.Store + store *coremetricspg.Store pool *pgxpool.Pool worker *execution.Worker registry *runtimegateway.Registry @@ -51,7 +51,7 @@ func (s *coreMetricsSource) Sample(ctx context.Context) coremetrics.Sample { if s.registry != nil { devices = s.registry.Devices() } - counts, err := s.store.ReadCoreExecutionSnapshot(ctx, time.Now(), devices) + counts, err := s.store.ReadExecutionSnapshot(ctx, time.Now(), devices) if err != nil { sample.Healthy = false } else { @@ -61,7 +61,7 @@ func (s *coreMetricsSource) Sample(ctx context.Context) coremetrics.Sample { sample.WaitingForDaemon = metricPtr(counts.WaitingForDaemon) } } - size, err := s.store.ReadCoreDatabaseSize(ctx) + size, err := s.store.ReadDatabaseSize(ctx) if err != nil { sample.Healthy = false } else { @@ -71,15 +71,7 @@ func (s *coreMetricsSource) Sample(ctx context.Context) coremetrics.Sample { return sample } func (s *coreMetricsSource) History(ctx context.Context, start, end time.Time, step time.Duration) (coremetrics.History, error) { - value, err := s.store.ReadCoreExecutionHistory(ctx, start, end, step) - if err != nil { - return coremetrics.History{}, err - } - result := coremetrics.History{Interrupted: value.Interrupted, QueueWaitMS: coremetrics.Latency{P50: value.QueueWaitMS.P50, P95: value.QueueWaitMS.P95}, Buckets: map[time.Time]*float64{}} - for _, bucket := range value.Buckets { - result.Buckets[bucket.Start.UTC()] = bucket.P95MS - } - return result, nil + return s.store.ReadExecutionHistory(ctx, start, end, step) } func reportCleanupResult(metrics *coremetrics.Service, job string, count int64, err error) { diff --git a/services/core/cmd/server/http_routes_test.go b/services/core/cmd/server/http_routes_test.go index 2af0668e5..998788235 100644 --- a/services/core/cmd/server/http_routes_test.go +++ b/services/core/cmd/server/http_routes_test.go @@ -105,7 +105,7 @@ func daemonComposition(t testing.TB) http.Handler { Artifacts: struct{ api.Artifacts }{}, ArtifactsReader: struct{ api.ArtifactsReader }{}, SessionAdmin: struct{ api.SessionAdmin }{}, Environments: struct{ api.Environments }{}, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, ExecutorConnections: struct{ api.ExecutorConnections }{}, - Admin: struct{ api.Admin }{}, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, Metrics: struct{ api.Metrics }{}, + Admin: struct{ api.Admin }{}, AdminRuntimeTargets: struct{ api.AdminRuntimeTargets }{}, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, Metrics: struct{ api.Metrics }{}, RuntimeObservations: struct{ api.RuntimeObservations }{}, RuntimeHistory: struct{ api.RuntimeHistory }{}, Execution: &api.Execution{ ExecutorURL: "wss://core.example/api/v1/agent-daemon/ws", diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index 044305694..ecbb7f82c 100644 --- a/services/core/cmd/server/main.go +++ b/services/core/cmd/server/main.go @@ -45,6 +45,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/nativeinstaller" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/agentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/coremetricspg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/filepg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/modelconfigurationpg" @@ -198,7 +199,7 @@ func run() error { if err != nil { return err } - metricsSource := &coreMetricsSource{store: executionStore, pool: pool} + metricsSource := &coreMetricsSource{store: coremetricspg.New(units), pool: pool} metrics := coremetrics.New(processStartedAt, buildRevision, metricsSource) auditRetention, err := writeAuditRetention() if err != nil { @@ -233,7 +234,7 @@ func run() error { if err != nil { return err } - history, err := runtimeHistory(ctx, executionStore, public != "") + history, err := runtimeHistory(ctx, units, public != "") if err != nil { return err } @@ -290,7 +291,7 @@ func run() error { if err != nil { return err } - daemonHandler, registry, err = runtime.NewGateway(sessionStore, sessionService, executionStore, executorURL) + daemonHandler, registry, err = runtime.NewGateway(sessionStore, sessionService, sessionStore, executorURL) if err != nil { return err } @@ -412,7 +413,7 @@ func run() error { ArtifactsReader: sessionStore, SessionAdmin: executionStore, Environments: sessionService, EnvironmentsReader: sessionStore, ExecutorConnections: executorConnections{sessions: sessionStore, registry: registry}, - Admin: executionStore, AdminAudit: auditStore, WriteAudit: auditStore, Metrics: metrics, + Admin: executionStore, AdminRuntimeTargets: sessionStore, AdminAudit: auditStore, WriteAudit: auditStore, Metrics: metrics, RuntimeObservations: observationService, RuntimeHistory: historyService, } if worker != nil { diff --git a/services/core/cmd/server/runtime_history.go b/services/core/cmd/server/runtime_history.go index 308e84a55..3ce6bf015 100644 --- a/services/core/cmd/server/runtime_history.go +++ b/services/core/cmd/server/runtime_history.go @@ -12,11 +12,11 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/coremetrics" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimehistory" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimehistory/postgresreader" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs/otlpexporter" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "golang.org/x/net/http/httpguts" ) @@ -52,7 +52,7 @@ type runtimeHistoryExporter interface { Close(context.Context) error } -func runtimeHistory(ctx context.Context, coreStore *store.Store, executionEnabled bool) (runtimeHistorySetup, error) { +func runtimeHistory(ctx context.Context, units *pgunit.Pool, executionEnabled bool) (runtimeHistorySetup, error) { config, err := loadRuntimeHistoryConfig() if err != nil { return runtimeHistorySetup{}, err @@ -70,7 +70,7 @@ func runtimeHistory(ctx context.Context, coreStore *store.Store, executionEnable Metrics: []runtimehistory.Metric{runtimehistory.MetricCPU, runtimehistory.MetricMemory, runtimehistory.MetricTokens}, } timeout := time.Duration(config.TimeoutSeconds) * time.Second - backend, err := postgresreader.New(coreStore, postgresreader.Config{Capabilities: capabilities, QueryTimeout: timeout}) + backend, err := postgresreader.New(units, postgresreader.Config{Capabilities: capabilities, QueryTimeout: timeout}) if err != nil { return runtimeHistorySetup{}, err } diff --git a/services/core/cmd/server/runtime_history_test.go b/services/core/cmd/server/runtime_history_test.go index 4cf172842..c640f2c92 100644 --- a/services/core/cmd/server/runtime_history_test.go +++ b/services/core/cmd/server/runtime_history_test.go @@ -2,8 +2,8 @@ package main import ( "context" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "os" "path/filepath" "strings" @@ -19,7 +19,7 @@ func (panicHistoryExporter) Close(context.Context) error func TestRuntimeHistoryUsesCoreDatabaseByDefault(t *testing.T) { t.Setenv("OAC_HISTORY_SETTINGS_FILE", "") for _, enabled := range []bool{true, false} { - setup, err := runtimeHistory(t.Context(), store.New(nil), enabled) + setup, err := runtimeHistory(t.Context(), pgunit.NewPool(nil), enabled) if err != nil { t.Fatal(err) } @@ -45,7 +45,7 @@ func TestRuntimeHistoryOptionalExportAndSamplingConfiguration(t *testing.T) { t.Fatal(err) } t.Setenv("OAC_HISTORY_SETTINGS_FILE", file) - setup, err := runtimeHistory(t.Context(), store.New(nil), true) + setup, err := runtimeHistory(t.Context(), pgunit.NewPool(nil), true) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/api/admin_resources.go b/services/core/internal/api/admin_resources.go index fbddfcf92..79ad09fe7 100644 --- a/services/core/internal/api/admin_resources.go +++ b/services/core/internal/api/admin_resources.go @@ -13,11 +13,9 @@ import ( // Only Core-key-authenticated project resource handlers receive it. type adminTenantContextKey struct{} -// Admin reads the administrator's cross-Project views: the asset summary and -// the Sessions whose Runtime is observed. +// Admin reads the administrator's asset summary of a Project. type Admin interface { ReadAdminSummary(context.Context, string, store.AdminSummaryFilter, func(sessions.Session, *string) error) (store.AdminAssetCounts, error) - ListAdminRuntimeTargets(context.Context, []string, string, int, bool) (store.AdminRuntimeTargetPage, error) } func (h *Handler) adminResourceScope(next http.Handler) http.Handler { diff --git a/services/core/internal/api/admin_runtime.go b/services/core/internal/api/admin_runtime.go index 393e65581..c0ff4716e 100644 --- a/services/core/internal/api/admin_runtime.go +++ b/services/core/internal/api/admin_runtime.go @@ -7,8 +7,15 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) +// AdminRuntimeTargets lists the Sessions whose Runtime the administrator +// observes across managed Projects. +type AdminRuntimeTargets interface { + ListAdminRuntimeTargets(context.Context, []string, string, int, bool) (sessions.AdminRuntimeTargetPage, error) +} + type AdminRuntimeObservation struct { ProjectID string `json:"project_id"` Observation AdminRuntimeObservationDetail `json:"observation"` @@ -72,9 +79,9 @@ func (h *Handler) adminRuntimeObservations(w http.ResponseWriter, r *http.Reques break } } - page, err := h.Admin.ListAdminRuntimeTargets(ctx, tenants, options.after, options.limit, options.ascending) + page, err := h.AdminRuntimeTargets.ListAdminRuntimeTargets(ctx, tenants, options.after, options.limit, options.ascending) if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } sessions := make([]runtimeobs.SessionIdentity, len(page.Data)) diff --git a/services/core/internal/api/admin_runtime_test.go b/services/core/internal/api/admin_runtime_test.go index 4507c074a..9df4f2579 100644 --- a/services/core/internal/api/admin_runtime_test.go +++ b/services/core/internal/api/admin_runtime_test.go @@ -11,19 +11,19 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" ) type adminRuntimeTargets struct { - page store.AdminRuntimeTargetPage + page sessions.AdminRuntimeTargetPage tenants []string after string limit int ascending bool } -func (s *adminRuntimeTargets) ListAdminRuntimeTargets(_ context.Context, tenants []string, after string, limit int, ascending bool) (store.AdminRuntimeTargetPage, error) { +func (s *adminRuntimeTargets) ListAdminRuntimeTargets(_ context.Context, tenants []string, after string, limit int, ascending bool) (sessions.AdminRuntimeTargetPage, error) { s.tenants, s.after, s.limit, s.ascending = tenants, after, limit, ascending return s.page, nil } @@ -32,14 +32,14 @@ const adminRuntimeObservationsPath = "/core/v1/sandbox/runtime-observations" // adminRuntimeFixture serves the administrator observation list over one // Project per tenant and the given Session targets. -func adminRuntimeFixture(t *testing.T, catalog []projects.Project, targets []store.AdminRuntimeTarget, service RuntimeObservations) (http.Handler, *adminRuntimeTargets) { +func adminRuntimeFixture(t *testing.T, catalog []projects.Project, targets []sessions.AdminRuntimeTarget, service RuntimeObservations) (http.Handler, *adminRuntimeTargets) { t.Helper() deps, fakes := managementFakes(t, callerBinding()) fakes.projectsReader.listProjects = func(context.Context, projects.ListQuery) (projects.Page, error) { return projects.Page{Data: catalog}, nil } - management := &adminRuntimeTargets{page: store.AdminRuntimeTargetPage{Data: targets, HasMore: true}} - fakes.admin.listAdminRuntimeTargets = management.ListAdminRuntimeTargets + management := &adminRuntimeTargets{page: sessions.AdminRuntimeTargetPage{Data: targets, HasMore: true}} + fakes.adminRuntimeTargets.listAdminRuntimeTargets = management.ListAdminRuntimeTargets observeWith(service)(&deps, fakes) return newTestHandler(t, deps), management } @@ -73,12 +73,12 @@ func unsupportedObservation(session string, at time.Time) runtimeobs.Observation func TestAdminRuntimeObservationListKeepsTargetOrderAndProjects(t *testing.T) { now := time.Date(2026, 9, 22, 8, 0, 0, 0, time.UTC) catalog := []projects.Project{{ID: uuid.NewString(), TenantID: uuid.NewString()}, {ID: uuid.NewString(), TenantID: uuid.NewString()}} - var targets []store.AdminRuntimeTarget + var targets []sessions.AdminRuntimeTarget for index := range 3 { - targets = append(targets, store.AdminRuntimeTarget{SessionID: uuid.NewString(), TenantID: catalog[index%2].TenantID}) + targets = append(targets, sessions.AdminRuntimeTarget{SessionID: uuid.NewString(), TenantID: catalog[index%2].TenantID}) } service := runtimeObservationServiceFunc(func(_ context.Context, tenant, session string) (runtimeobs.Observation, error) { - if !slices.Contains(targets, store.AdminRuntimeTarget{SessionID: session, TenantID: tenant}) { + if !slices.Contains(targets, sessions.AdminRuntimeTarget{SessionID: session, TenantID: tenant}) { return runtimeobs.Observation{}, errors.New("unexpected tenant") } return unsupportedObservation(session, now), nil @@ -107,9 +107,9 @@ func TestAdminRuntimeObservationListKeepsTargetOrderAndProjects(t *testing.T) { func TestAdminRuntimeObservationListBoundsCollection(t *testing.T) { now := time.Date(2026, 9, 22, 8, 0, 0, 0, time.UTC) project := projects.Project{ID: uuid.NewString(), TenantID: uuid.NewString()} - var targets []store.AdminRuntimeTarget + var targets []sessions.AdminRuntimeTarget for range 20 { - targets = append(targets, store.AdminRuntimeTarget{SessionID: uuid.NewString(), TenantID: project.TenantID}) + targets = append(targets, sessions.AdminRuntimeTarget{SessionID: uuid.NewString(), TenantID: project.TenantID}) } var options runtimeobs.PageOptions service := runtimeObservationPageRecorder{options: &options, runtimeObservationServiceFunc: func(_ context.Context, _, session string) (runtimeobs.Observation, error) { @@ -134,7 +134,7 @@ func TestAdminRuntimeObservationListRejectsWholePageOnIntegrityFailure(t *testin } return unsupportedObservation(session, now), nil }) - handler, _ := adminRuntimeFixture(t, []projects.Project{project}, []store.AdminRuntimeTarget{{SessionID: valid, TenantID: project.TenantID}, {SessionID: invalid, TenantID: project.TenantID}}, service) + handler, _ := adminRuntimeFixture(t, []projects.Project{project}, []sessions.AdminRuntimeTarget{{SessionID: valid, TenantID: project.TenantID}, {SessionID: invalid, TenantID: project.TenantID}}, service) response := runtimeObservationRequest(handler, adminRuntimeObservationsPath+"?limit=2") var envelope struct { diff --git a/services/core/internal/api/dependencies.go b/services/core/internal/api/dependencies.go index d1e25c875..c810fa9c6 100644 --- a/services/core/internal/api/dependencies.go +++ b/services/core/internal/api/dependencies.go @@ -55,6 +55,7 @@ type Dependencies struct { EnvironmentsReader EnvironmentsReader ExecutorConnections ExecutorConnections Admin Admin + AdminRuntimeTargets AdminRuntimeTargets AdminAudit AdminAudit WriteAudit WriteAudit Metrics Metrics @@ -141,7 +142,7 @@ func (d Dependencies) validate() error { field{"Artifacts", d.Artifacts}, field{"ArtifactsReader", d.ArtifactsReader}, field{"SessionAdmin", d.SessionAdmin}, field{"Environments", d.Environments}, field{"EnvironmentsReader", d.EnvironmentsReader}, - field{"ExecutorConnections", d.ExecutorConnections}, field{"Admin", d.Admin}, field{"AdminAudit", d.AdminAudit}, field{"WriteAudit", d.WriteAudit}, + field{"ExecutorConnections", d.ExecutorConnections}, field{"Admin", d.Admin}, field{"AdminRuntimeTargets", d.AdminRuntimeTargets}, field{"AdminAudit", d.AdminAudit}, field{"WriteAudit", d.WriteAudit}, field{"Metrics", d.Metrics}, field{"RuntimeObservations", d.RuntimeObservations}, field{"RuntimeHistory", d.RuntimeHistory}, ); err != nil { return err diff --git a/services/core/internal/api/dependencies_test.go b/services/core/internal/api/dependencies_test.go index 41722853c..698629487 100644 --- a/services/core/internal/api/dependencies_test.go +++ b/services/core/internal/api/dependencies_test.go @@ -41,6 +41,7 @@ type testFakes struct { environmentsReader *fakeEnvironmentsReader executorConnections *fakeExecutorConnections admin *fakeAdmin + adminRuntimeTargets *fakeAdminRuntimeTargets adminAudit *fakeAdminAudit writeAudit *fakeWriteAudit metrics *fakeMetrics @@ -84,7 +85,7 @@ func testDependencies(t testing.TB) (Dependencies, *testFakes) { artifacts: &fakeArtifacts{t: t}, artifactsReader: &fakeArtifactsReader{t: t}, sessionAdmin: &fakeSessionAdmin{t: t}, environments: &fakeEnvironments{t: t}, environmentsReader: &fakeEnvironmentsReader{t: t}, executorConnections: &fakeExecutorConnections{t: t}, - admin: &fakeAdmin{t: t}, adminAudit: &fakeAdminAudit{t: t}, writeAudit: &fakeWriteAudit{t: t}, metrics: &fakeMetrics{t: t}, + admin: &fakeAdmin{t: t}, adminRuntimeTargets: &fakeAdminRuntimeTargets{t: t}, adminAudit: &fakeAdminAudit{t: t}, writeAudit: &fakeWriteAudit{t: t}, metrics: &fakeMetrics{t: t}, runtimeObservations: &fakeRuntimeObservations{t: t}, runtimeHistory: &fakeRuntimeHistory{t: t}, installationBindings: &fakeInstallationBindings{t: t}, sessionAdmission: &fakeSessionAdmission{t: t}, inputAdmission: &fakeInputAdmission{t: t}, @@ -111,7 +112,7 @@ func testDependencies(t testing.TB) (Dependencies, *testFakes) { Artifacts: f.artifacts, ArtifactsReader: f.artifactsReader, SessionAdmin: f.sessionAdmin, - Environments: f.environments, EnvironmentsReader: f.environmentsReader, ExecutorConnections: f.executorConnections, Admin: f.admin, AdminAudit: f.adminAudit, WriteAudit: f.writeAudit, + Environments: f.environments, EnvironmentsReader: f.environmentsReader, ExecutorConnections: f.executorConnections, Admin: f.admin, AdminRuntimeTargets: f.adminRuntimeTargets, AdminAudit: f.adminAudit, WriteAudit: f.writeAudit, Metrics: f.metrics, RuntimeObservations: f.runtimeObservations, RuntimeHistory: f.runtimeHistory, }, f } diff --git a/services/core/internal/api/fakes_test.go b/services/core/internal/api/fakes_test.go index e51df824c..bd21b0a51 100644 --- a/services/core/internal/api/fakes_test.go +++ b/services/core/internal/api/fakes_test.go @@ -42,9 +42,8 @@ func unexpectedCall(t testing.TB, method string) { } type fakeAdmin struct { - t testing.TB - readAdminSummary func(context.Context, string, store.AdminSummaryFilter, func(sessions.Session, *string) error) (store.AdminAssetCounts, error) - listAdminRuntimeTargets func(context.Context, []string, string, int, bool) (store.AdminRuntimeTargetPage, error) + t testing.TB + readAdminSummary func(context.Context, string, store.AdminSummaryFilter, func(sessions.Session, *string) error) (store.AdminAssetCounts, error) } func (f *fakeAdmin) ReadAdminSummary(a0 context.Context, a1 string, a2 store.AdminSummaryFilter, a3 func(sessions.Session, *string) error) (store.AdminAssetCounts, error) { @@ -54,7 +53,12 @@ func (f *fakeAdmin) ReadAdminSummary(a0 context.Context, a1 string, a2 store.Adm return f.readAdminSummary(a0, a1, a2, a3) } -func (f *fakeAdmin) ListAdminRuntimeTargets(a0 context.Context, a1 []string, a2 string, a3 int, a4 bool) (store.AdminRuntimeTargetPage, error) { +type fakeAdminRuntimeTargets struct { + t testing.TB + listAdminRuntimeTargets func(context.Context, []string, string, int, bool) (sessions.AdminRuntimeTargetPage, error) +} + +func (f *fakeAdminRuntimeTargets) ListAdminRuntimeTargets(a0 context.Context, a1 []string, a2 string, a3 int, a4 bool) (sessions.AdminRuntimeTargetPage, error) { if f.listAdminRuntimeTargets == nil { unexpectedCall(f.t, "ListAdminRuntimeTargets") } diff --git a/services/core/internal/coremetrics/service.go b/services/core/internal/coremetrics/service.go index 87fb79d1e..b5433c329 100644 --- a/services/core/internal/coremetrics/service.go +++ b/services/core/internal/coremetrics/service.go @@ -2,7 +2,6 @@ package coremetrics import ( "context" - "errors" "regexp" "runtime" "sync" @@ -120,7 +119,7 @@ func Window(now time.Time, name string) (Range, error) { case "7d": duration, step = retention, 2*time.Hour default: - return Range{}, errors.New("invalid Core metrics range") + return Range{}, ErrInvalidRange } end := now.UTC().Truncate(step) return Range{Start: end.Add(-duration), End: end, ResolutionSeconds: int64(step / time.Second)}, nil diff --git a/services/core/internal/coremetrics/types.go b/services/core/internal/coremetrics/types.go index 3cc899d16..f4c132e5c 100644 --- a/services/core/internal/coremetrics/types.go +++ b/services/core/internal/coremetrics/types.go @@ -3,9 +3,15 @@ package coremetrics import ( "context" + "errors" "time" ) +// ErrInvalidRange reports a Core metrics range that is not one of the supported +// windows, or a history interval that is unbounded or not aligned to its +// resolution. +var ErrInvalidRange = errors.New("invalid Core metrics range") + type Latency struct { P50 *float64 `json:"p50" extensions:"x-nullable"` P95 *float64 `json:"p95" extensions:"x-nullable"` @@ -103,6 +109,15 @@ type Live struct { Pool Pool Scheduler Job } + +// ExecutionSnapshot counts the deployment's persisted root Turns. +type ExecutionSnapshot struct { + QueuedTurns, WaitingForDaemon, InProgressTurns int64 + OldestQueuedSeconds *float64 +} + +// History is the root Turn history of a range. Buckets maps each bucket's UTC +// start to its queue wait p95 in milliseconds, nil without observations. type History struct { Interrupted int64 QueueWaitMS Latency diff --git a/services/core/internal/execution/archive_cancellation_cleanup_test.go b/services/core/internal/execution/archive_cancellation_cleanup_test.go index 6618cb37f..21a4a3916 100644 --- a/services/core/internal/execution/archive_cancellation_cleanup_test.go +++ b/services/core/internal/execution/archive_cancellation_cleanup_test.go @@ -116,7 +116,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" credentials, heartbeat := testSessions(t, pool, testCredentialCipher(t)) - handler, liveRegistry, err := runtime.NewGateway(credentials, heartbeat, s, wsURL) + handler, liveRegistry, err := runtime.NewGateway(credentials, heartbeat, credentials, wsURL) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/persistence/postgres/coremetricspg/store.go b/services/core/internal/persistence/postgres/coremetricspg/store.go new file mode 100644 index 000000000..50ec3fad8 --- /dev/null +++ b/services/core/internal/persistence/postgres/coremetricspg/store.go @@ -0,0 +1,107 @@ +// Package coremetricspg reads the Core metrics that PostgreSQL holds: the +// deployment's root Turn queue and its history, and the database size. +package coremetricspg + +import ( + "context" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/coremetrics" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" +) + +// Store reads Core metrics from Core's database. +type Store struct { + units *pgunit.Pool +} + +// New builds the Core metrics reader on units. +func New(units *pgunit.Pool) *Store { return &Store{units: units} } + +// ReadExecutionSnapshot counts persisted root Turns across the deployment. +// WaitingForDaemon is the queued subset whose Session binding is absent from the +// supplied live registry IDs; callers must distinguish an unavailable registry +// from an observed empty one before using this method. Age is in seconds, and is +// nil when the queue is empty. No Session deletion filter hides operational state. +func (s *Store) ReadExecutionSnapshot(ctx context.Context, now time.Time, connectedDeviceIDs []string) (coremetrics.ExecutionSnapshot, error) { + ids := make([]pgtype.UUID, 0, len(connectedDeviceIDs)) + for _, value := range connectedDeviceIDs { + id, err := pgunit.ParseID(value) + if err != nil { + return coremetrics.ExecutionSnapshot{}, err + } + ids = append(ids, id) + } + row, err := s.units.Queries().CoreExecutionSnapshot(ctx, sqlc.CoreExecutionSnapshotParams{ + ConnectedDeviceIds: ids, ObservedAt: pgtype.Timestamptz{Time: now, Valid: true}, + }) + if err != nil { + return coremetrics.ExecutionSnapshot{}, err + } + result := coremetrics.ExecutionSnapshot{QueuedTurns: row.QueuedTurns, WaitingForDaemon: row.WaitingForDaemon, InProgressTurns: row.InProgressTurns} + if row.QueuedTurns > 0 { + result.OldestQueuedSeconds = &row.OldestQueuedSeconds + } + return result, nil +} + +// ReadExecutionHistory reads a bounded, repeatable read-only snapshot of root +// Turn history. The interval is [start,end), with complete epoch-aligned buckets +// keyed by their UTC start. Interruptions use failed Turns' completed_at and +// execution_interrupted error code. Queue waits use started_at-created_at in +// milliseconds, grouped by started_at, including retained history of deleted +// Sessions. Counts of an empty interval are zero; percentile values without +// observations are nil. Read errors invalidate the entire result and must not be +// presented as measured zeros. +func (s *Store) ReadExecutionHistory(ctx context.Context, start, end time.Time, resolution time.Duration) (coremetrics.History, error) { + span := end.Sub(start) + if resolution < time.Second || resolution%time.Second != 0 || span <= 0 || span > 7*24*time.Hour || + span%resolution != 0 || span/resolution > 1008 || start.Nanosecond() != 0 || end.Nanosecond() != 0 || + start.Unix()%int64(resolution/time.Second) != 0 || end.Unix()%int64(resolution/time.Second) != 0 { + return coremetrics.History{}, coremetrics.ErrInvalidRange + } + buckets := make([]time.Time, int(span/resolution)) + result := coremetrics.History{Buckets: make(map[time.Time]*float64, len(buckets))} + for i := range buckets { + buckets[i] = start.Add(time.Duration(i) * resolution).UTC() + result.Buckets[buckets[i]] = nil + } + first, last := pgtype.Timestamptz{Time: start, Valid: true}, pgtype.Timestamptz{Time: end, Valid: true} + err := s.units.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + var err error + result.Interrupted, err = q.CoreInterruptedTurns(ctx, sqlc.CoreInterruptedTurnsParams{RangeStart: first, RangeEnd: last}) + if err != nil { + return err + } + wait, err := q.CoreQueueWaitSummary(ctx, sqlc.CoreQueueWaitSummaryParams{RangeStart: first, RangeEnd: last}) + if err != nil { + return err + } + if wait.Samples > 0 { + result.QueueWaitMS = coremetrics.Latency{P50: &wait.P50Ms, P95: &wait.P95Ms} + } + rows, err := q.CoreQueueWaitBuckets(ctx, sqlc.CoreQueueWaitBucketsParams{RangeStart: first, RangeEnd: last, ResolutionSeconds: int32(resolution / time.Second)}) + if err != nil { + return err + } + for _, row := range rows { + if row.Samples > 0 { + result.Buckets[buckets[row.BucketNumber]] = &row.P95Ms + } + } + return nil + }) + if err != nil { + return coremetrics.History{}, err + } + return result, nil +} + +// ReadDatabaseSize measures the current PostgreSQL database in bytes. +func (s *Store) ReadDatabaseSize(ctx context.Context) (int64, error) { + return s.units.Queries().CoreDatabaseSize(ctx) +} diff --git a/services/core/internal/store/core_metrics_test.go b/services/core/internal/persistence/postgres/coremetricspg/store_test.go similarity index 79% rename from services/core/internal/store/core_metrics_test.go rename to services/core/internal/persistence/postgres/coremetricspg/store_test.go index 9fb2f6c2a..db5d8edba 100644 --- a/services/core/internal/store/core_metrics_test.go +++ b/services/core/internal/persistence/postgres/coremetricspg/store_test.go @@ -1,4 +1,4 @@ -package store +package coremetricspg import ( "context" @@ -8,10 +8,11 @@ import ( "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/coremetrics" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) func coreMetricsSession(t *testing.T, pool *pgxpool.Pool, deleted bool) string { @@ -42,9 +43,10 @@ func coreMetricsTurn(t *testing.T, pool *pgxpool.Pool, session, status, code str } func TestCoreMetricsSnapshot(t *testing.T) { - s, pool := testStore(t) + pool := pgtest.Open(t) + s := New(pgunit.NewPool(pool)) now := time.Now().UTC() - baseline, err := s.ReadCoreExecutionSnapshot(t.Context(), now, []string{}) + baseline, err := s.ReadExecutionSnapshot(t.Context(), now, []string{}) if err != nil { t.Fatal(err) } @@ -66,7 +68,7 @@ func TestCoreMetricsSnapshot(t *testing.T) { } } } - got, err := s.ReadCoreExecutionSnapshot(t.Context(), now, []string{connected}) + got, err := s.ReadExecutionSnapshot(t.Context(), now, []string{connected}) if err != nil { t.Fatal(err) } @@ -76,17 +78,18 @@ func TestCoreMetricsSnapshot(t *testing.T) { if got.OldestQueuedSeconds == nil || math.Abs(*got.OldestQueuedSeconds-now.Sub(oldest).Seconds()) > 0.001 { t.Fatalf("oldest age: %v", got.OldestQueuedSeconds) } - disconnected, err := s.ReadCoreExecutionSnapshot(t.Context(), now, []string{}) + disconnected, err := s.ReadExecutionSnapshot(t.Context(), now, []string{}) if err != nil || disconnected.WaitingForDaemon != got.WaitingForDaemon+1 { t.Fatalf("registry disconnect: %+v, %v", disconnected, err) } - if _, err := s.ReadCoreExecutionSnapshot(t.Context(), now, []string{"invalid-device-id"}); err == nil { + if _, err := s.ReadExecutionSnapshot(t.Context(), now, []string{"invalid-device-id"}); err == nil { t.Fatal("invalid registry ID accepted") } } func TestCoreMetricsHistory(t *testing.T) { - s, pool := testStore(t) + pool := pgtest.Open(t) + s := New(pgunit.NewPool(pool)) start := time.Date(2001, 1, 1, 0, 0, 0, 0, time.UTC) end := start.Add(3 * time.Minute) // Boundary samples exercise inclusive start, exclusive end and exact buckets. @@ -109,7 +112,7 @@ func TestCoreMetricsHistory(t *testing.T) { } coreMetricsTurn(t, pool, coreMetricsSession(t, pool, false), status, code, start, nil, &completed) } - got, err := s.ReadCoreExecutionHistory(t.Context(), start, end, time.Minute) + got, err := s.ReadExecutionHistory(t.Context(), start, end, time.Minute) if err != nil { t.Fatal(err) } @@ -124,33 +127,33 @@ func TestCoreMetricsHistory(t *testing.T) { } check("range p50", got.QueueWaitMS.P50, 1000) check("range p95", got.QueueWaitMS.P95, 2800) - check("first bucket p95", got.Buckets[0].P95MS, 950) - check("second bucket p95", got.Buckets[1].P95MS, 3000) - if got.Buckets[2].P95MS != nil || !got.Buckets[2].Start.Equal(start.Add(2*time.Minute)) { + check("first bucket p95", got.Buckets[start], 950) + check("second bucket p95", got.Buckets[start.Add(time.Minute)], 3000) + if p95, ok := got.Buckets[start.Add(2*time.Minute)]; !ok || p95 != nil { t.Fatal("missing bucket must have aligned start and unknown percentile") } - empty, err := s.ReadCoreExecutionHistory(t.Context(), end.Add(time.Hour), end.Add(2*time.Hour), time.Minute) + empty, err := s.ReadExecutionHistory(t.Context(), end.Add(time.Hour), end.Add(2*time.Hour), time.Minute) if err != nil || empty.Interrupted != 0 || empty.QueueWaitMS.P50 != nil || empty.QueueWaitMS.P95 != nil || len(empty.Buckets) != 60 { t.Fatalf("empty history: %+v, %v", empty, err) } - for _, bucket := range empty.Buckets { - if bucket.P95MS != nil { + for _, p95 := range empty.Buckets { + if p95 != nil { t.Fatal("empty bucket invented zero percentile") } } - size, err := s.ReadCoreDatabaseSize(t.Context()) + size, err := s.ReadDatabaseSize(t.Context()) if err != nil || size <= 0 { t.Fatalf("database size: %d, %v", size, err) } ctx, cancel := context.WithCancel(t.Context()) cancel() - if _, err := s.ReadCoreExecutionSnapshot(ctx, start, nil); err == nil { + if _, err := s.ReadExecutionSnapshot(ctx, start, nil); err == nil { t.Fatal("snapshot read error hidden") } - if result, err := s.ReadCoreExecutionHistory(ctx, start, end, time.Minute); err == nil || result.Buckets != nil { + if result, err := s.ReadExecutionHistory(ctx, start, end, time.Minute); err == nil || result.Buckets != nil { t.Fatal("history read error hidden or partial data returned") } - if _, err := s.ReadCoreDatabaseSize(ctx); err == nil { + if _, err := s.ReadDatabaseSize(ctx); err == nil { t.Fatal("database size read error hidden") } } @@ -171,7 +174,7 @@ func TestCoreMetricsHistoryBounds(t *testing.T) { {start, start.Add(time.Hour + time.Second), time.Minute}, {start.Add(time.Nanosecond), start.Add(time.Hour + time.Nanosecond), time.Minute}, } { - if _, err := s.ReadCoreExecutionHistory(t.Context(), tc.start, tc.end, tc.step); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := s.ReadExecutionHistory(t.Context(), tc.start, tc.end, tc.step); !errors.Is(err, coremetrics.ErrInvalidRange) { t.Fatalf("unbounded or unaligned range accepted: %+v, %v", tc, err) } } diff --git a/services/core/internal/persistence/postgres/sessionpg/admin.go b/services/core/internal/persistence/postgres/sessionpg/admin.go new file mode 100644 index 000000000..3f49c20be --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/admin.go @@ -0,0 +1,55 @@ +package sessionpg + +import ( + "context" + "errors" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" +) + +func (s *Store) ListAdminRuntimeTargets(ctx context.Context, tenantIDs []string, after string, limit int, ascending bool) (sessions.AdminRuntimeTargetPage, error) { + page := sessions.AdminRuntimeTargetPage{Data: []sessions.AdminRuntimeTarget{}} + if limit < 1 || limit > 100 { + return page, sessions.ErrInvalidInput + } + tenants := make([]pgtype.UUID, 0, len(tenantIDs)) + for _, value := range tenantIDs { + id, err := parseID(value) + if err != nil { + return page, err + } + tenants = append(tenants, id) + } + q := s.units.Queries() + params := sqlc.AdminRuntimeTargetsParams{TenantIds: tenants, Ascending: ascending, AfterID: pgtype.UUID{Valid: true}, PageLimit: int32(limit + 1)} + if after != "" { + var err error + params.AfterID, err = parseID(after) + if err != nil { + return page, sessions.ErrNotFound + } + params.AfterTime, err = q.AdminRuntimeCursor(ctx, sqlc.AdminRuntimeCursorParams{ID: params.AfterID, TenantIds: tenants}) + if errors.Is(err, pgx.ErrNoRows) { + return page, sessions.ErrNotFound + } + if err != nil { + return page, err + } + } + rows, err := q.AdminRuntimeTargets(ctx, params) + if err != nil { + return page, err + } + page.HasMore = len(rows) > limit + if page.HasMore { + rows = rows[:limit] + } + for _, row := range rows { + page.Data = append(page.Data, sessions.AdminRuntimeTarget{SessionID: uuid.UUID(row.ID.Bytes).String(), TenantID: uuid.UUID(row.TenantID.Bytes).String()}) + } + return page, nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/admin_test.go b/services/core/internal/persistence/postgres/sessionpg/admin_test.go new file mode 100644 index 000000000..9d1e34609 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/admin_test.go @@ -0,0 +1,63 @@ +package sessionpg + +import ( + "errors" + "reflect" + "testing" + "time" + + "github.com/google/uuid" + + "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 TestListAdminRuntimeTargetsPagesLiveSessionsOfTenants(t *testing.T) { + pool := pgtest.Open(t) + store := New(pgunit.NewPool(pool), nil) + first, second, foreign := uuid.NewString(), uuid.NewString(), uuid.NewString() + created := time.Now().UTC().Add(-time.Hour) + session := func(tenant string, offset time.Duration, deleted bool) sessions.AdminRuntimeTarget { + id := uuid.NewString() + exec(t, pool, `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash, created_at, deleted_at) + VALUES ($1, $2, 'codex', $5, 'hash', $3, CASE WHEN $4::boolean THEN clock_timestamp() END)`, id, tenant, created.Add(offset), deleted, id) + return sessions.AdminRuntimeTarget{SessionID: id, TenantID: tenant} + } + a := session(first, 0, false) + b := session(second, time.Second, false) + c := session(first, 2*time.Second, false) + deleted := session(first, 3*time.Second, true) + other := session(foreign, 4*time.Second, false) + tenants := []string{first, second} + + page, err := store.ListAdminRuntimeTargets(t.Context(), tenants, "", 2, true) + if err != nil || !page.HasMore || !reflect.DeepEqual(page.Data, []sessions.AdminRuntimeTarget{a, b}) { + t.Fatalf("first ascending page = %+v, %v", page, err) + } + page, err = store.ListAdminRuntimeTargets(t.Context(), tenants, b.SessionID, 2, true) + if err != nil || page.HasMore || !reflect.DeepEqual(page.Data, []sessions.AdminRuntimeTarget{c}) { + t.Fatalf("second ascending page = %+v, %v", page, err) + } + page, err = store.ListAdminRuntimeTargets(t.Context(), tenants, "", 100, false) + if err != nil || page.HasMore || !reflect.DeepEqual(page.Data, []sessions.AdminRuntimeTarget{c, b, a}) { + t.Fatalf("descending page = %+v, %v", page, err) + } + page, err = store.ListAdminRuntimeTargets(t.Context(), []string{}, "", 20, false) + if err != nil || page.HasMore || page.Data == nil || len(page.Data) != 0 { + t.Fatalf("page without tenants = %+v, %v", page, err) + } + for _, after := range []string{deleted.SessionID, other.SessionID, "not-a-session"} { + if _, err := store.ListAdminRuntimeTargets(t.Context(), tenants, after, 20, false); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("cursor %q: %v", after, err) + } + } + for _, limit := range []int{0, 101} { + if _, err := store.ListAdminRuntimeTargets(t.Context(), tenants, "", limit, false); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatalf("limit %d: %v", limit, err) + } + } + if _, err := store.ListAdminRuntimeTargets(t.Context(), []string{"not-a-tenant"}, "", 20, false); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatalf("malformed tenant: %v", err) + } +} diff --git a/services/core/internal/persistence/postgres/sessionpg/devices.go b/services/core/internal/persistence/postgres/sessionpg/devices.go index 30b04e7f1..0db12ccb9 100644 --- a/services/core/internal/persistence/postgres/sessionpg/devices.go +++ b/services/core/internal/persistence/postgres/sessionpg/devices.go @@ -76,6 +76,29 @@ func (s *Store) GetDeviceCredential(ctx context.Context, device string) (runtime }, true, nil } +// ArchivedCancellationReceipt is a read-only exception for the exact already +// authenticated delivery. The ordinary credential view remains revoked; this +// cannot authorize bootstrap, reconnect, dispatch, workspace access or renewal. +// A marker records that archive caused the first revocation; timestamps alone +// cannot distinguish an earlier ordinary cancel/revoke followed by archive. +func (s *Store) ArchivedCancellationReceipt(ctx context.Context, device, credentialHash string, runIDs []string) (runtimedevice.ArchivedCancellationReceipt, error) { + if len(runIDs) == 0 || credentialHash == "" { + return runtimedevice.ArchivedCancellationReceipt{}, nil + } + id, err := parseID(device) + if err != nil { + return runtimedevice.ArchivedCancellationReceipt{}, err + } + row, err := s.units.Queries().GetArchivedCancellationReceipt(ctx, sqlc.GetArchivedCancellationReceiptParams{DeviceID: id, CredentialHash: pgtype.Text{String: credentialHash, Valid: true}, RunIds: runIDs, LimitSeconds: int32(runtimedevice.ArchivedCancellationReceiptLimit.Seconds())}) + if errors.Is(err, pgx.ErrNoRows) { + return runtimedevice.ArchivedCancellationReceipt{}, nil + } + if err != nil { + return runtimedevice.ArchivedCancellationReceipt{}, err + } + return runtimedevice.ArchivedCancellationReceipt{RunID: uuid.UUID(row.ID.Bytes).String(), Deadline: row.CancelRequestedAt.Time.Add(runtimedevice.ArchivedCancellationReceiptLimit)}, nil +} + func (s *Store) ListEnrolledRuntimeBindings(ctx context.Context) ([]sessions.EnrolledRuntimeBinding, error) { rows, err := s.units.Queries().ListEnrolledRuntimeBindings(ctx) if err != nil { diff --git a/services/core/internal/runtimehistory/postgresreader/reader.go b/services/core/internal/runtimehistory/postgresreader/reader.go index 6c8672482..1011ec484 100644 --- a/services/core/internal/runtimehistory/postgresreader/reader.go +++ b/services/core/internal/runtimehistory/postgresreader/reader.go @@ -5,6 +5,7 @@ import ( "errors" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimehistory" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" "github.com/google/uuid" @@ -15,6 +16,8 @@ import ( const maximumRawSamples = 20_000 const retention = 7 * 24 * time.Hour +// sampleStore is the sample storage the Reader runs on: samples in Core, a +// fake in the unit tests. type sampleStore interface { InsertRuntimeHistorySample(context.Context, runtimeobs.ExportRecord) error ListRuntimeHistorySamples(context.Context, string, string, string, int64, int64, int32) ([]runtimeobs.ExportRecord, error) @@ -33,8 +36,16 @@ type Reader struct { now func() time.Time } -func New(store sampleStore, config Config) (*Reader, error) { - if store == nil || config.QueryTimeout <= 0 || config.Capabilities.Validate() != nil || config.Capabilities.Retention != retention || config.Capabilities.MaximumRange > 24*time.Hour { +// New builds the Runtime history backend on Core's database. +func New(units *pgunit.Pool, config Config) (*Reader, error) { + if units == nil { + return nil, errors.New("invalid Runtime history PostgreSQL configuration") + } + return newReader(samples{units: units}, config) +} + +func newReader(store sampleStore, config Config) (*Reader, error) { + if config.QueryTimeout <= 0 || config.Capabilities.Validate() != nil || config.Capabilities.Retention != retention || config.Capabilities.MaximumRange > 24*time.Hour { return nil, errors.New("invalid Runtime history PostgreSQL configuration") } capabilities := config.Capabilities diff --git a/services/core/internal/runtimehistory/postgresreader/reader_test.go b/services/core/internal/runtimehistory/postgresreader/reader_test.go index e2fc8b28a..7d75176da 100644 --- a/services/core/internal/runtimehistory/postgresreader/reader_test.go +++ b/services/core/internal/runtimehistory/postgresreader/reader_test.go @@ -65,7 +65,7 @@ func testReader(t *testing.T, s *fakeStore, now time.Time) *Reader { t.Helper() caps := testCapabilities() caps.Metrics = append(caps.Metrics, runtimehistory.MetricTokens) - r, err := New(s, Config{Capabilities: caps, QueryTimeout: time.Second}) + r, err := newReader(s, Config{Capabilities: caps, QueryTimeout: time.Second}) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/runtime_history.go b/services/core/internal/runtimehistory/postgresreader/samples.go similarity index 72% rename from services/core/internal/store/runtime_history.go rename to services/core/internal/runtimehistory/postgresreader/samples.go index 912db93ef..fb0302972 100644 --- a/services/core/internal/store/runtime_history.go +++ b/services/core/internal/runtimehistory/postgresreader/samples.go @@ -1,28 +1,34 @@ -package store +package postgresreader import ( "context" "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgtype" ) +// samples stores Runtime history samples in Core's database. +type samples struct { + units *pgunit.Pool +} + // InsertRuntimeHistorySample copies a sanitized periodic observation. Public // Session/Turn Usage remains the accounting authority; these measured counters // are only sampled chart values. Deleted or mismatched owners cannot create orphan rows. -func (s *Store) InsertRuntimeHistorySample(ctx context.Context, record runtimeobs.ExportRecord) error { - tenant, err := parseID(record.TenantID) +func (s samples) InsertRuntimeHistorySample(ctx context.Context, record runtimeobs.ExportRecord) error { + tenant, err := pgunit.ParseID(record.TenantID) if err != nil { return err } - session, err := parseID(record.SessionID) + session, err := pgunit.ParseID(record.SessionID) if err != nil { return err } - environment, err := parseID(record.EnvironmentID) + environment, err := pgunit.ParseID(record.EnvironmentID) if err != nil { return err } @@ -31,7 +37,7 @@ func (s *Store) InsertRuntimeHistorySample(ctx context.Context, record runtimeob ResolvedAtNs: record.ResolvedAt.UnixNano(), ProviderType: record.ProviderType, Status: string(record.Status), } if record.AllocationID != "" { - params.AllocationID, err = parseID(record.AllocationID) + params.AllocationID, err = pgunit.ParseID(record.AllocationID) if err != nil { return err } @@ -51,23 +57,23 @@ func (s *Store) InsertRuntimeHistorySample(ctx context.Context, record runtimeob params.InputTokens = historyInteger(&usage.InputTokens) params.OutputTokens = historyInteger(&usage.OutputTokens) } - return s.queries.InsertRuntimeHistorySample(ctx, params) + return s.units.Queries().InsertRuntimeHistorySample(ctx, params) } -func (s *Store) ListRuntimeHistorySamples(ctx context.Context, tenantID, sessionID, environmentID string, startNS, endNS int64, limit int32) ([]runtimeobs.ExportRecord, error) { - tenant, err := parseID(tenantID) +func (s samples) ListRuntimeHistorySamples(ctx context.Context, tenantID, sessionID, environmentID string, startNS, endNS int64, limit int32) ([]runtimeobs.ExportRecord, error) { + tenant, err := pgunit.ParseID(tenantID) if err != nil { return nil, err } - session, err := parseID(sessionID) + session, err := pgunit.ParseID(sessionID) if err != nil { return nil, err } - environment, err := parseID(environmentID) + environment, err := pgunit.ParseID(environmentID) if err != nil { return nil, err } - rows, err := s.queries.ListRuntimeHistorySamples(ctx, sqlc.ListRuntimeHistorySamplesParams{ + rows, err := s.units.Queries().ListRuntimeHistorySamples(ctx, sqlc.ListRuntimeHistorySamplesParams{ TenantID: tenant, SessionID: session, EnvironmentID: environment, StartNs: startNS, EndNs: endNS, RowLimit: limit, }) if err != nil { @@ -103,12 +109,15 @@ func (s *Store) ListRuntimeHistorySamples(ctx context.Context, tenantID, session return records, nil } -func (s *Store) PruneRuntimeHistorySamples(ctx context.Context, beforeNS int64) (int64, error) { - count, err := s.queries.PruneRuntimeHistorySamples(ctx, beforeNS) +// PruneRuntimeHistorySamples deletes one batch of expired Runtime samples and +// one batch of expired node-host samples, which share the retention. +func (s samples) PruneRuntimeHistorySamples(ctx context.Context, beforeNS int64) (int64, error) { + queries := s.units.Queries() + count, err := queries.PruneRuntimeHistorySamples(ctx, beforeNS) if err != nil { return count, err } - nodes, err := s.queries.PruneNodeHostHistory(ctx, pgtype.Timestamptz{Time: time.Unix(0, beforeNS), Valid: true}) + nodes, err := queries.PruneNodeHostHistory(ctx, pgtype.Timestamptz{Time: time.Unix(0, beforeNS), Valid: true}) return count + nodes, err } diff --git a/services/core/internal/store/runtime_history_acceptance_test.go b/services/core/internal/runtimehistory/postgresreader/samples_test.go similarity index 76% rename from services/core/internal/store/runtime_history_acceptance_test.go rename to services/core/internal/runtimehistory/postgresreader/samples_test.go index fb496e743..9ba566f2c 100644 --- a/services/core/internal/store/runtime_history_acceptance_test.go +++ b/services/core/internal/runtimehistory/postgresreader/samples_test.go @@ -1,16 +1,14 @@ -package store_test +package postgresreader import ( "context" - "encoding/json" "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimehistory" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimehistory/postgresreader" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" ) @@ -20,26 +18,26 @@ func historyCapabilities() runtimehistory.Capabilities { Retention: 7 * 24 * time.Hour, MinimumStep: 30 * time.Second, MaximumRange: 24 * time.Hour, MaximumPoints: 1000, MaximumSeries: 64, MaximumTotalPoints: 10000, Metrics: []runtimehistory.Metric{runtimehistory.MetricCPU, runtimehistory.MetricMemory, runtimehistory.MetricTokens}} } -func historyBackend(t *testing.T, s *store.Store) *postgresreader.Reader { +func historyBackend(t *testing.T, pool *pgxpool.Pool) *Reader { t.Helper() - r, err := postgresreader.New(s, postgresreader.Config{Capabilities: historyCapabilities(), QueryTimeout: 5 * time.Second}) + r, err := New(pgunit.NewPool(pool), Config{Capabilities: historyCapabilities(), QueryTimeout: 5 * time.Second}) if err != nil { t.Fatal(err) } return r } -func historyOwner(t *testing.T, s *store.Store, pool *pgxpool.Pool) runtimehistory.Scope { + +// historyOwner seeds a Session of a fresh tenant with its Environment. +func historyOwner(t *testing.T, pool *pgxpool.Pool) runtimehistory.Scope { t.Helper() - tenant := uuid.NewString() - session, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"model":"fixture-model"},"environment":{"type":"openai_hosted","workspace_directory":"/workspace","capability_directories":[]}}`)}) - if err != nil { + scope := runtimehistory.Scope{TenantID: uuid.NewString(), SessionID: uuid.NewString(), EnvironmentID: uuid.NewString()} + if _, err := pool.Exec(t.Context(), `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash) VALUES ($1, $2, 'codex', 'key', 'hash')`, scope.SessionID, scope.TenantID); err != nil { t.Fatal(err) } - environment, err := sessionReads(pool).GetSessionEnvironment(t.Context(), tenant, session.ID) - if err != nil { + if _, err := pool.Exec(t.Context(), `INSERT INTO environments(id, session_id) VALUES ($1, $2)`, scope.EnvironmentID, scope.SessionID); err != nil { t.Fatal(err) } - return runtimehistory.Scope{TenantID: tenant, SessionID: session.ID, EnvironmentID: environment.ID} + return scope } func historyRecord(scope runtimehistory.Scope, allocation string, started, at time.Time, cpu float64, input uint64) runtimeobs.ExportRecord { capacity := 2.0 @@ -56,10 +54,10 @@ func historyQuery(scope runtimehistory.Scope, start, end time.Time, points int) } func TestPostgresRuntimeHistoryAcceptance(t *testing.T) { - s, pool := store.NewTestStore(t) - scope := historyOwner(t, s, pool) - foreign := historyOwner(t, s, pool) - reader := historyBackend(t, s) + pool := pgtest.Open(t) + scope := historyOwner(t, pool) + foreign := historyOwner(t, pool) + reader := historyBackend(t, pool) end := time.Now().UTC().Truncate(time.Second) start := end.Add(-3 * time.Minute) started := start.Add(-time.Minute + 123*time.Nanosecond) @@ -101,7 +99,7 @@ func TestPostgresRuntimeHistoryAcceptance(t *testing.T) { t.Fatal(err) } defer freshPool.Close() - reader = historyBackend(t, store.New(freshPool)) + reader = historyBackend(t, freshPool) q := historyQuery(scope, start, end, 6) result, err := reader.Query(t.Context(), q) if err != nil { @@ -129,17 +127,17 @@ func TestPostgresRuntimeHistoryAcceptance(t *testing.T) { t.Fatal("scope isolation failed", got, err) } } - // Chart snapshots do not alter canonical Session Usage. - session, err := s.GetSession(t.Context(), scope.TenantID, scope.SessionID) - if err != nil || len(session.Usage) != 0 && string(session.Usage) != "null" { - t.Fatal("telemetry became accounting authority", string(session.Usage), err) + // Chart snapshots do not alter canonical Session Usage, which Turns carry. + var usage int + if err := pool.QueryRow(t.Context(), `SELECT count(*) FROM turns WHERE session_id = $1 AND token_usage IS NOT NULL`, scope.SessionID).Scan(&usage); err != nil || usage != 0 { + t.Fatal("telemetry became accounting authority", usage, err) } } func TestPostgresRuntimeHistoryDenseReadAndBoundedRetention(t *testing.T) { - s, pool := store.NewTestStore(t) - scope := historyOwner(t, s, pool) - reader := historyBackend(t, s) + pool := pgtest.Open(t) + scope := historyOwner(t, pool) + reader := historyBackend(t, pool) end := time.Now().UTC().Truncate(time.Second) start := end.Add(-24 * time.Hour) started := start.Add(-time.Hour + 987*time.Nanosecond) @@ -171,7 +169,7 @@ func TestPostgresRuntimeHistoryDenseReadAndBoundedRetention(t *testing.T) { } // Expired rows remain invisible before physical pruning runs. old := historyRecord(scope, base.AllocationID, started.Add(-8*24*time.Hour), start.Add(-8*24*time.Hour), 0, 1) - if err := s.InsertRuntimeHistorySample(t.Context(), old); err != nil { + if err := (samples{units: pgunit.NewPool(pool)}).InsertRuntimeHistorySample(t.Context(), old); err != nil { t.Fatal(err) } _, err = pool.Exec(t.Context(), `INSERT INTO runtime_history_samples @@ -206,9 +204,9 @@ func TestPostgresRuntimeHistoryDenseReadAndBoundedRetention(t *testing.T) { } func TestPostgresRuntimeHistoryKeepsProviderReportedUtilization(t *testing.T) { - s, pool := store.NewTestStore(t) - scope := historyOwner(t, s, pool) - reader := historyBackend(t, s) + pool := pgtest.Open(t) + scope := historyOwner(t, pool) + reader := historyBackend(t, pool) allocation := uuid.NewString() start := time.Now().UTC().Truncate(time.Minute).Add(-10 * time.Minute) started := start.Add(-time.Minute) @@ -228,3 +226,25 @@ func TestPostgresRuntimeHistoryKeepsProviderReportedUtilization(t *testing.T) { t.Fatalf("bucket utilization = %v, want the mean of reported ratios", ratio) } } + +// Retention also expires node-host samples and never deletes nodes. +func TestNodeHostHistoryRetentionKeepsNode(t *testing.T) { + pool := pgtest.Open(t) + node := uuid.NewString() + if _, err := pool.Exec(t.Context(), `INSERT INTO runtime_nodes(id, installation_id, name, backend_fingerprint, credential_sha256, max_active, max_retained) VALUES ($1, $2, 'history', repeat('a', 64), repeat('b', 64), 1, 1)`, node, uuid.NewString()); err != nil { + t.Fatal(err) + } + now := time.Now().UTC() + for _, at := range []time.Time{now.Add(-8 * 24 * time.Hour), now.Add(-2 * time.Minute)} { + if _, err := pool.Exec(t.Context(), `INSERT INTO node_host_history_samples(node_id, observed_at) VALUES ($1, $2)`, node, at); err != nil { + t.Fatal(err) + } + } + if _, err := historyBackend(t, pool).Prune(t.Context()); err != nil { + t.Fatal(err) + } + var expired, kept, nodes int + if err := pool.QueryRow(t.Context(), `SELECT count(*) FILTER (WHERE observed_at < $2), count(*) FILTER (WHERE observed_at >= $2), (SELECT count(*) FROM runtime_nodes WHERE id = $1) FROM node_host_history_samples WHERE node_id = $1`, node, now.Add(-7*24*time.Hour)).Scan(&expired, &kept, &nodes); err != nil || expired != 0 || kept != 1 || nodes != 1 { + t.Fatal("retention kept expired samples or removed recent ones or the node", expired, kept, nodes, err) + } +} diff --git a/services/core/internal/sandbox/providers/configuration_flow_test.go b/services/core/internal/sandbox/providers/configuration_flow_test.go index 541b350f1..6b1ac416f 100644 --- a/services/core/internal/sandbox/providers/configuration_flow_test.go +++ b/services/core/internal/sandbox/providers/configuration_flow_test.go @@ -142,7 +142,7 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { Artifacts: struct{ api.Artifacts }{}, ArtifactsReader: struct{ api.ArtifactsReader }{}, SessionAdmin: s, - Environments: struct{ api.Environments }{}, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, Admin: s, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, + Environments: struct{ api.Environments }{}, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, Admin: s, AdminRuntimeTargets: struct{ api.AdminRuntimeTargets }{}, 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/admin.go b/services/core/internal/sessions/admin.go new file mode 100644 index 000000000..c68310671 --- /dev/null +++ b/services/core/internal/sessions/admin.go @@ -0,0 +1,23 @@ +package sessions + +import "context" + +// AdminRuntimeTarget names a live Session whose Runtime an administrator +// observes, with its tenant. +type AdminRuntimeTarget struct{ SessionID, TenantID string } + +// AdminRuntimeTargetPage is one page of AdminRuntimeTargets. +type AdminRuntimeTargetPage struct { + Data []AdminRuntimeTarget + HasMore bool +} + +// AdminReader reads the administrator's cross-Project Session views. +type AdminReader interface { + // ListAdminRuntimeTargets pages the live Sessions of the tenants by + // creation time, then ID, ascending or descending, after the Session + // after. A limit outside 1 to 100 or a malformed tenant is + // ErrInvalidInput; an after Session that is malformed, deleted or + // outside the tenants is ErrNotFound. + ListAdminRuntimeTargets(ctx context.Context, tenants []string, after string, limit int, ascending bool) (AdminRuntimeTargetPage, error) +} diff --git a/services/core/internal/sessions/devices.go b/services/core/internal/sessions/devices.go index b419c3983..417263183 100644 --- a/services/core/internal/sessions/devices.go +++ b/services/core/internal/sessions/devices.go @@ -65,6 +65,10 @@ type DeviceReader interface { // authenticates a device with, and reports whether the device still has // authority. GetDeviceCredential(ctx context.Context, device string) (runtimedevice.Credential, bool, error) + // ArchivedCancellationReceipt reads the receipt window that archiving a + // Session leaves the device's exact authenticated delivery of one of + // runIDs; without one it is the zero receipt. + ArchivedCancellationReceipt(ctx context.Context, device, credentialHash string, runIDs []string) (runtimedevice.ArchivedCancellationReceipt, error) // ListEnrolledRuntimeBindings lists the enrolled user-managed Runtimes of // live Environments. ListEnrolledRuntimeBindings(ctx context.Context) ([]EnrolledRuntimeBinding, error) diff --git a/services/core/internal/sessions/reader.go b/services/core/internal/sessions/reader.go index 5fe5d1e34..fe80400aa 100644 --- a/services/core/internal/sessions/reader.go +++ b/services/core/internal/sessions/reader.go @@ -1,8 +1,10 @@ package sessions -// Reader reads a Session's resources, Environment and devices, one family per -// line. Callers use it directly; no use case forwards a read. +// Reader reads a Session's resources, Environment and devices, and the +// administrator's cross-Project Session views, one family per line. Callers +// use it directly; no use case forwards a read. type Reader interface { + AdminReader ArtifactReader DeviceReader EnvironmentReader diff --git a/services/core/internal/store/admin_summary.go b/services/core/internal/store/admin_summary.go index 43364b8e6..97fffd512 100644 --- a/services/core/internal/store/admin_summary.go +++ b/services/core/internal/store/admin_summary.go @@ -2,13 +2,11 @@ package store import ( "context" - "errors" "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" ) @@ -75,54 +73,6 @@ func (s *Store) ReadAdminSummary(ctx context.Context, tenantID string, filter Ad return counts, err } -type AdminRuntimeTarget struct{ SessionID, TenantID string } -type AdminRuntimeTargetPage struct { - Data []AdminRuntimeTarget - HasMore bool -} - -func (s *Store) ListAdminRuntimeTargets(ctx context.Context, tenantIDs []string, after string, limit int, ascending bool) (AdminRuntimeTargetPage, error) { - page := AdminRuntimeTargetPage{Data: []AdminRuntimeTarget{}} - if limit < 1 || limit > 100 { - return page, sessions.ErrInvalidInput - } - tenants := make([]pgtype.UUID, 0, len(tenantIDs)) - for _, value := range tenantIDs { - id, err := parseID(value) - if err != nil { - return page, err - } - tenants = append(tenants, id) - } - params := sqlc.AdminRuntimeTargetsParams{TenantIds: tenants, Ascending: ascending, AfterID: pgtype.UUID{Valid: true}, PageLimit: int32(limit + 1)} - if after != "" { - var err error - params.AfterID, err = parseID(after) - if err != nil { - return page, sessions.ErrNotFound - } - params.AfterTime, err = s.queries.AdminRuntimeCursor(ctx, sqlc.AdminRuntimeCursorParams{ID: params.AfterID, TenantIds: tenants}) - if errors.Is(err, pgx.ErrNoRows) { - return page, sessions.ErrNotFound - } - if err != nil { - return page, err - } - } - rows, err := s.queries.AdminRuntimeTargets(ctx, params) - if err != nil { - return page, err - } - page.HasMore = len(rows) > limit - if page.HasMore { - rows = rows[:limit] - } - for _, row := range rows { - page.Data = append(page.Data, AdminRuntimeTarget{SessionID: uuid.UUID(row.ID.Bytes).String(), TenantID: uuid.UUID(row.TenantID.Bytes).String()}) - } - return page, nil -} - func summaryTimestamp(value *time.Time) pgtype.Timestamptz { if value == nil { return pgtype.Timestamptz{} diff --git a/services/core/internal/store/archive_cancellation_test.go b/services/core/internal/store/archive_cancellation_test.go index 6a34f1f5f..3434dd0cb 100644 --- a/services/core/internal/store/archive_cancellation_test.go +++ b/services/core/internal/store/archive_cancellation_test.go @@ -87,7 +87,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { } server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), s, wsURL) + handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), fixtureSessionStore(db), wsURL) if err != nil { t.Fatal(err) } @@ -199,11 +199,11 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { if dialErr == nil || response == nil || response.StatusCode != http.StatusUnauthorized { t.Fatal("revoked Runtime reconnected") } - drain, err := s.ArchivedCancellationReceipt(t.Context(), owner.DeviceID, secret, nil) + drain, err := fixtureSessionStore(db).ArchivedCancellationReceipt(t.Context(), owner.DeviceID, secret, nil) if err != nil || drain.RunID != "" { t.Fatal("unowned delivery got receipt permission", drain, err) } - drain, err = s.ArchivedCancellationReceipt(t.Context(), owner.DeviceID, runtimedevice.HashCredential(secret), []string{input.TurnID}) + drain, err = fixtureSessionStore(db).ArchivedCancellationReceipt(t.Context(), owner.DeviceID, runtimedevice.HashCredential(secret), []string{input.TurnID}) if err != nil || (drain.RunID == input.TurnID) == strings.Contains(scenario, "revoke") { t.Fatal("archive revocation causality lost", drain, err) } diff --git a/services/core/internal/store/core_metrics.go b/services/core/internal/store/core_metrics.go deleted file mode 100644 index e6e30dcbb..000000000 --- a/services/core/internal/store/core_metrics.go +++ /dev/null @@ -1,111 +0,0 @@ -package store - -import ( - "context" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -type CoreExecutionSnapshot struct { - QueuedTurns, WaitingForDaemon, InProgressTurns int64 - OldestQueuedSeconds *float64 -} - -type CoreQueueWait struct{ P50, P95 *float64 } - -type CoreQueueWaitBucket struct { - Start time.Time - P95MS *float64 -} - -type CoreExecutionHistory struct { - Interrupted int64 - QueueWaitMS CoreQueueWait - Buckets []CoreQueueWaitBucket -} - -// ReadCoreExecutionSnapshot counts persisted root Turns across the deployment. -// WaitingForDaemon is the queued subset whose Session binding is absent from the -// supplied live registry IDs; callers must distinguish an unavailable registry -// from an observed empty one before using this method. Age is in seconds, and is -// nil when the queue is empty. No Session deletion filter hides operational state. -func (s *Store) ReadCoreExecutionSnapshot(ctx context.Context, now time.Time, connectedDeviceIDs []string) (CoreExecutionSnapshot, error) { - ids := make([]pgtype.UUID, 0, len(connectedDeviceIDs)) - for _, value := range connectedDeviceIDs { - id, err := parseID(value) - if err != nil { - return CoreExecutionSnapshot{}, err - } - ids = append(ids, id) - } - row, err := s.queries.CoreExecutionSnapshot(ctx, sqlc.CoreExecutionSnapshotParams{ - ConnectedDeviceIds: ids, ObservedAt: pgtype.Timestamptz{Time: now, Valid: true}, - }) - if err != nil { - return CoreExecutionSnapshot{}, err - } - result := CoreExecutionSnapshot{QueuedTurns: row.QueuedTurns, WaitingForDaemon: row.WaitingForDaemon, InProgressTurns: row.InProgressTurns} - if row.QueuedTurns > 0 { - result.OldestQueuedSeconds = &row.OldestQueuedSeconds - } - return result, nil -} - -// ReadCoreExecutionHistory reads a bounded, repeatable read-only snapshot of root -// Turn history. The interval is [start,end), with complete epoch-aligned buckets. -// Interruptions use failed Turns' completed_at and execution_interrupted error -// code. Queue waits use started_at-created_at in milliseconds, grouped by -// started_at, including retained history of deleted Sessions. Counts of an empty -// interval are zero; percentile values without observations are nil. Read errors -// invalidate the entire result and must not be presented as measured zeros. -func (s *Store) ReadCoreExecutionHistory(ctx context.Context, start, end time.Time, resolution time.Duration) (CoreExecutionHistory, error) { - span := end.Sub(start) - if resolution < time.Second || resolution%time.Second != 0 || span <= 0 || span > 7*24*time.Hour || - span%resolution != 0 || span/resolution > 1008 || start.Nanosecond() != 0 || end.Nanosecond() != 0 || - start.Unix()%int64(resolution/time.Second) != 0 || end.Unix()%int64(resolution/time.Second) != 0 { - return CoreExecutionHistory{}, sessions.ErrInvalidInput - } - result := CoreExecutionHistory{Buckets: make([]CoreQueueWaitBucket, int(span/resolution))} - for i := range result.Buckets { - result.Buckets[i].Start = start.Add(time.Duration(i) * resolution).UTC() - } - first, last := pgtype.Timestamptz{Time: start, Valid: true}, pgtype.Timestamptz{Time: end, Valid: true} - err := s.pooled.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - var err error - result.Interrupted, err = q.CoreInterruptedTurns(ctx, sqlc.CoreInterruptedTurnsParams{RangeStart: first, RangeEnd: last}) - if err != nil { - return err - } - wait, err := q.CoreQueueWaitSummary(ctx, sqlc.CoreQueueWaitSummaryParams{RangeStart: first, RangeEnd: last}) - if err != nil { - return err - } - if wait.Samples > 0 { - result.QueueWaitMS = CoreQueueWait{P50: &wait.P50Ms, P95: &wait.P95Ms} - } - rows, err := q.CoreQueueWaitBuckets(ctx, sqlc.CoreQueueWaitBucketsParams{RangeStart: first, RangeEnd: last, ResolutionSeconds: int32(resolution / time.Second)}) - if err != nil { - return err - } - for _, row := range rows { - if row.Samples > 0 { - result.Buckets[row.BucketNumber].P95MS = &row.P95Ms - } - } - return nil - }) - if err != nil { - return CoreExecutionHistory{}, err - } - return result, nil -} - -// ReadCoreDatabaseSize measures the current PostgreSQL database in bytes. -func (s *Store) ReadCoreDatabaseSize(ctx context.Context) (int64, error) { - return s.queries.CoreDatabaseSize(ctx) -} diff --git a/services/core/internal/store/devices_test.go b/services/core/internal/store/devices_test.go index 1e1d8f3e7..eacff7428 100644 --- a/services/core/internal/store/devices_test.go +++ b/services/core/internal/store/devices_test.go @@ -124,7 +124,7 @@ func TestStandaloneGatewayUsesExecutionCredentials(t *testing.T) { _, foreignSecret := registerTestDevice(t, s, uuid.NewString()) server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - handler, registry, err := runtime.NewGateway(sessionAdapter(s), sessionService(t, s), s, wsURL) + handler, registry, err := runtime.NewGateway(sessionAdapter(s), sessionService(t, s), sessionAdapter(s), wsURL) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/dispatch_test.go b/services/core/internal/store/dispatch_test.go index cb8fbe2a5..112b42b3a 100644 --- a/services/core/internal/store/dispatch_test.go +++ b/services/core/internal/store/dispatch_test.go @@ -88,7 +88,7 @@ func newDispatchHarnessForSession(t *testing.T, configuration []byte, local bool } server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - server.Config.Handler, h.registry, err = runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), s, wsURL) + server.Config.Handler, h.registry, err = runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), fixtureSessionStore(db), wsURL) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/node_host_history_test.go b/services/core/internal/store/node_host_history_test.go deleted file mode 100644 index 43bdcc818..000000000 --- a/services/core/internal/store/node_host_history_test.go +++ /dev/null @@ -1,23 +0,0 @@ -package store - -import ( - "testing" - "time" -) - -// Retention uses the Runtime history cleanup operation and never deletes nodes. -func TestNodeHostHistoryRetentionKeepsNode(t *testing.T) { - s, _, d := managerFixture(t, 2, 8) - now := time.Now().UTC() - for _, at := range []time.Time{now.Add(-8 * 24 * time.Hour), now.Add(-2 * time.Minute)} { - if _, err := s.pool.Exec(t.Context(), "INSERT INTO node_host_history_samples(node_id, observed_at) VALUES($1,$2)", d.LocalNodeID, at); err != nil { - t.Fatal(err) - } - } - if n, err := s.PruneRuntimeHistorySamples(t.Context(), now.Add(-7*24*time.Hour).UnixNano()); err != nil || n != 1 { - t.Fatal(n, err) - } - if _, err := deploymentService(t, s).NodeDetail(t.Context(), d.LocalNodeID, "1h"); err != nil { - t.Fatal(err) - } -} diff --git a/services/core/internal/store/public_handler_fixture_test.go b/services/core/internal/store/public_handler_fixture_test.go index 82da7f7e6..b04fa9f73 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: sessionService, EnvironmentsReader: sessionStore, Admin: s, AdminAudit: audit, WriteAudit: audit, + SessionAdmin: s, Environments: sessionService, EnvironmentsReader: sessionStore, Admin: s, AdminRuntimeTargets: sessionStore, AdminAudit: audit, WriteAudit: audit, ExecutorConnections: strict, Metrics: strict, RuntimeObservations: strict, RuntimeHistory: strict, } for _, c := range configure { diff --git a/services/core/internal/store/runtime_cancellation.go b/services/core/internal/store/runtime_cancellation.go deleted file mode 100644 index 6f14db569..000000000 --- a/services/core/internal/store/runtime_cancellation.go +++ /dev/null @@ -1,35 +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/runtimedevice" - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -// ArchivedCancellationReceipt is a read-only exception for the exact already -// authenticated delivery. The ordinary credential view remains revoked; this -// cannot authorize bootstrap, reconnect, dispatch, workspace access or renewal. -// A marker records that archive caused the first revocation; timestamps alone -// cannot distinguish an earlier ordinary cancel/revoke followed by archive. -func (s *Store) ArchivedCancellationReceipt(ctx context.Context, deviceID, credentialHash string, runIDs []string) (runtimedevice.ArchivedCancellationReceipt, error) { - if len(runIDs) == 0 || credentialHash == "" { - return runtimedevice.ArchivedCancellationReceipt{}, nil - } - id, err := parseID(deviceID) - if err != nil { - return runtimedevice.ArchivedCancellationReceipt{}, err - } - row, err := s.queries.GetArchivedCancellationReceipt(ctx, sqlc.GetArchivedCancellationReceiptParams{DeviceID: id, CredentialHash: pgtype.Text{String: credentialHash, Valid: true}, RunIds: runIDs, LimitSeconds: int32(runtimedevice.ArchivedCancellationReceiptLimit.Seconds())}) - if errors.Is(err, pgx.ErrNoRows) { - return runtimedevice.ArchivedCancellationReceipt{}, nil - } - if err != nil { - return runtimedevice.ArchivedCancellationReceipt{}, err - } - return runtimedevice.ArchivedCancellationReceipt{RunID: uuid.UUID(row.ID.Bytes).String(), Deadline: row.CancelRequestedAt.Time.Add(runtimedevice.ArchivedCancellationReceiptLimit)}, nil -} diff --git a/services/core/internal/store/runtime_connection_test.go b/services/core/internal/store/runtime_connection_test.go index 088ff0eea..de89b8fe9 100644 --- a/services/core/internal/store/runtime_connection_test.go +++ b/services/core/internal/store/runtime_connection_test.go @@ -25,7 +25,7 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { tenant, session, environment := managedSession(t, s, db) server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), s, wsURL) + handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), fixtureSessionStore(db), wsURL) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/runtime_enrollment_connection_test.go b/services/core/internal/store/runtime_enrollment_connection_test.go index aea56cc8c..fde167d43 100644 --- a/services/core/internal/store/runtime_enrollment_connection_test.go +++ b/services/core/internal/store/runtime_enrollment_connection_test.go @@ -47,7 +47,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { } server := httptest.NewUnstartedServer(nil) wsURL := "ws://" + server.Listener.Addr().String() + "/api/v1/agent-daemon/ws" - handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), s, wsURL) + handler, registry, err := runtime.NewGateway(fixtureSessionStore(db), fixtureSessionService(t, db), fixtureSessionStore(db), wsURL) if err != nil { t.Fatal(err) } From d0a19f94b0135188a8f5bce60a24527b3db189a7 Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 01:22:00 +0800 Subject: [PATCH 2/6] Move Turn transitions and completion into sessions (#434) * Move Turn transitions and completion into sessions * Type Session model execution errors and run the Turn journal through WithTurns SessionModelExecution returns sessions.ErrNotFound for a Session without a frozen provider, credentialcrypto.ErrUnavailable without the credential key, and wraps the cipher error. AppendTurnEvents runs through WithTurns, so WithTurnJournal and the JournalExecution port are gone. Tests drop the partial Reader fake, the shadowed ApplyTurnStatus fake and an unused store export, and the lease test moves Turns through the leased operations. --- services/core/IMPLEMENTATION.md | 6 +- services/core/cmd/server/main.go | 2 +- .../native_classification_integration_test.go | 8 +- .../session_diagnostics_public_compat_test.go | 28 +- .../internal/api/session_diagnostics_test.go | 8 +- .../archive_cancellation_cleanup_test.go | 10 +- services/core/internal/execution/artifacts.go | 4 +- services/core/internal/execution/delivery.go | 4 +- .../deployment_provider_observations_test.go | 11 +- .../core/internal/execution/dispatcher.go | 8 +- .../execution/executor_preparation.go | 2 +- services/core/internal/execution/functions.go | 7 +- .../internal/execution/model_execution.go | 6 +- .../execution/model_execution_test.go | 11 +- .../core/internal/execution/owner_test.go | 9 +- .../internal/execution/prepared_dispatch.go | 2 +- services/core/internal/execution/worker.go | 8 +- .../core/internal/execution/worker_device.go | 6 +- .../internal/execution/worker_schedule.go | 4 +- .../persistence/postgres/sessionpg/binding.go | 6 +- .../persistence/postgres/sessionpg/devices.go | 54 +++- .../postgres/sessionpg/execution_journal.go | 24 -- .../postgres/sessionpg/execution_turns.go | 80 +++++ .../sessionpg/execution_turns_test.go | 206 ++++++++++++ .../postgres/sessionpg/functions.go | 15 - .../postgres/sessionpg/model_execution.go | 54 ++++ .../postgres/sessionpg/turn_reads.go | 119 +++++++ .../providers/configuration_flow_test.go | 2 +- services/core/internal/sessions/devices.go | 8 + services/core/internal/sessions/doc.go | 12 +- services/core/internal/sessions/execution.go | 2 +- .../internal/sessions/execution_functions.go | 13 +- .../sessions/execution_functions_test.go | 24 +- .../internal/sessions/execution_journal.go | 30 -- .../core/internal/sessions/execution_turns.go | 306 ++++++++++++++++++ .../internal/sessions/execution_turns_test.go | 294 +++++++++++++++++ .../core/internal/sessions/functions_test.go | 6 - .../core/internal/sessions/journal_test.go | 2 +- .../core/internal/sessions/model_execution.go | 16 + services/core/internal/sessions/reader.go | 2 + .../internal/sessions/transaction_test.go | 32 +- services/core/internal/sessions/turn.go | 22 ++ .../store/admin_session_archive_test.go | 2 +- .../admin_session_archive_worker_http_test.go | 2 +- .../agent_execution_defaults_http_test.go | 4 +- .../store/archive_cancellation_test.go | 2 +- .../core/internal/store/artifact_lifecycle.go | 48 --- .../internal/store/claude_execution_test.go | 10 +- .../core/internal/store/claude_mcp_test.go | 5 +- .../internal/store/command_output_test.go | 4 +- .../deployment_model_providers_http_test.go | 8 +- services/core/internal/store/devices.go | 44 --- services/core/internal/store/dispatch_test.go | 15 +- .../store/environment_claim_worker_test.go | 5 +- .../store/environment_directory_test.go | 3 +- .../store/environment_expiry_dispatch_test.go | 2 +- .../store/environment_input_activity_test.go | 2 +- .../store/environment_input_claim_test.go | 2 +- .../core/internal/store/environment_inputs.go | 7 +- .../internal/store/environment_setup_test.go | 2 +- .../store/environment_steering_order_test.go | 6 +- .../store/environment_worker_helpers_test.go | 2 +- .../internal/store/execution_events_test.go | 2 +- .../core/internal/store/execution_test.go | 55 +++- services/core/internal/store/export_test.go | 9 + .../store/function_execution_native_test.go | 2 +- .../internal/store/function_execution_test.go | 4 +- .../store/function_images_native_test.go | 4 +- .../store/function_inputs_public_test.go | 6 +- .../internal/store/function_inputs_test.go | 4 +- .../store/function_public_native_test.go | 2 +- .../store/function_state_public_test.go | 4 +- .../internal/store/function_state_test.go | 4 +- .../store/function_stream_native_test.go | 2 +- .../internal/store/function_worker_test.go | 2 +- .../internal/store/harness_onboarding_test.go | 5 +- .../core/internal/store/input_batches_test.go | 6 +- .../store/input_conflicts_public_test.go | 4 +- .../store/item_order_migration_test.go | 2 +- .../core/internal/store/item_order_test.go | 6 +- .../core/internal/store/item_reads_test.go | 16 +- .../internal/store/list_cursor_public_test.go | 4 +- .../store/local_environment_devices_test.go | 2 +- .../store/local_environment_worker_test.go | 4 +- .../store/mcode_public_native_test.go | 9 +- .../mcp_credential_selection_public_test.go | 2 +- .../store/message_images_native_test.go | 6 +- .../store/model_protocol_native_test.go | 4 +- .../internal/store/native_environment_test.go | 4 +- .../internal/store/native_recovery_test.go | 4 +- .../internal/store/no_environment_test.go | 2 +- .../store/prepared_dispatch_failure_test.go | 2 +- .../internal/store/prepared_dispatch_test.go | 4 +- .../internal/store/public_execution_test.go | 6 +- .../store/public_handler_fixture_test.go | 2 +- .../runtime_capabilities_pending_test.go | 2 +- .../internal/store/runtime_enrollment_test.go | 2 +- .../internal/store/runtime_idle_clock_test.go | 16 +- .../internal/store/runtime_initialization.go | 25 -- .../store/runtime_initialization_test.go | 4 +- .../runtime_suspension_concurrency_test.go | 6 +- .../store/runtime_worker_recovery_test.go | 4 +- .../store/sandbox_deployment_switch_test.go | 4 +- services/core/internal/store/scheduling.go | 55 +--- .../store/self_hosted_cancel_public_test.go | 5 +- .../store/session_artifacts_public_test.go | 2 +- .../internal/store/session_artifacts_test.go | 6 +- .../store/session_creation_identity_test.go | 6 +- .../store/session_creation_stream_test.go | 2 +- .../store/session_deletion_execution_test.go | 7 +- .../session_deletion_lifecycle_public_test.go | 4 +- .../internal/store/session_deletion_test.go | 12 +- .../internal/store/session_diagnostics.go | 2 +- .../store/session_diagnostics_test.go | 4 +- .../internal/store/session_events_test.go | 4 +- .../internal/store/session_metadata_test.go | 4 +- .../internal/store/session_model_execution.go | 32 -- .../store/session_model_execution_test.go | 10 +- .../store/session_model_options_test.go | 4 +- .../store/structured_output_dispatch_test.go | 2 +- .../store/structured_output_native_test.go | 6 +- .../store/subagent_identities_test.go | 2 +- .../store/subagent_native_outputs_test.go | 4 +- .../internal/store/subagent_resources_test.go | 8 +- .../store/subagent_visibility_public_test.go | 7 +- .../store/token_usage_integration_test.go | 22 +- .../internal/store/tool_policy_native_test.go | 7 +- .../internal/store/tool_search_native_test.go | 6 +- .../core/internal/store/turn_completion.go | 105 ------ .../core/internal/store/turn_events_test.go | 10 +- services/core/internal/store/turn_inputs.go | 6 +- .../core/internal/store/turn_inputs_test.go | 15 +- .../internal/store/turn_journal_reads_test.go | 2 +- services/core/internal/store/turn_reads.go | 49 --- .../core/internal/store/turn_reads_test.go | 71 ---- services/core/internal/store/turns.go | 116 ------- services/core/internal/store/turns_test.go | 68 +--- .../unified_model_configuration_http_test.go | 3 +- .../internal/store/worker_capacity_test.go | 7 +- .../internal/store/worker_input_race_test.go | 2 +- .../internal/store/worker_lease_loss_test.go | 6 +- services/core/tests/fixtures/items.go | 23 +- services/core/tests/fixtures/main.go | 4 +- 143 files changed, 1615 insertions(+), 1000 deletions(-) delete mode 100644 services/core/internal/persistence/postgres/sessionpg/execution_journal.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/execution_turns.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/execution_turns_test.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/model_execution.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/turn_reads.go delete mode 100644 services/core/internal/sessions/execution_journal.go create mode 100644 services/core/internal/sessions/execution_turns.go create mode 100644 services/core/internal/sessions/execution_turns_test.go create mode 100644 services/core/internal/sessions/model_execution.go delete mode 100644 services/core/internal/store/artifact_lifecycle.go delete mode 100644 services/core/internal/store/devices.go delete mode 100644 services/core/internal/store/runtime_initialization.go delete mode 100644 services/core/internal/store/turn_completion.go delete mode 100644 services/core/internal/store/turn_reads.go delete mode 100644 services/core/internal/store/turn_reads_test.go delete mode 100644 services/core/internal/store/turns.go diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 29e210fc0..2254cb2de 100644 --- a/services/core/IMPLEMENTATION.md +++ b/services/core/IMPLEMENTATION.md @@ -16,9 +16,9 @@ Two errors are shared across domains, each with one `api` helper: `textvalue.Err Shared vocabulary has one owner each, and domains use it rather than copy it. `internal/environmentconfig` owns Environment setup, Skills, Plugins and initial files with their validation and public metadata; `Setup.Validate` checks requested configuration, where a Skill may be an unresolved reference, and `Setup.ValidateInstalled` checks frozen, installable configuration. `internal/skills` owns `ParseVersion`, the canonical positive decimal Skill version. `internal/metadata` owns the metadata rules: `Validate` for the pair, key and value limits and U+0000, `ValidateStorable` for U+0000 alone, and `Encode` with its 64 KiB bound. `internal/jsonobject` owns `Normalize`, the stable encoding of stored JSON objects that snapshots and retry identities compare. These packages import no persistence. -`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. +`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, `LoadRequiredActions` lists the calls a Turn waits on, and `TransitionTurn` moves a root Turn's status as its current status allows and settles a Turn that ends. `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 and Turn operation loads the Turn in one leased Session transaction, decides in `sessions`, then applies in `sessionpg`, and `AppendTurnEvents` records and projects a Turn's journal batch in one leased Session transaction. `internal/items` owns public Items: it projects observations, merges them into stored Items and builds the ordered events that report each Item change. Neither imports persistence. `internal/persistence/postgres/sessionpg` loads the facts those decisions read and applies them inside the caller's Session transaction, under the Session lock: it allocates event sequence positions, event IDs, Item positions, output indexes and Subagent IDs, writes the change journal, each Turn's execution journal and its counters, Items, Turn usage, Subagent bindings, child Turns and child Items, and Artifact settlement, and prunes the change journal. It decides nothing. `WithSession` runs a Session operation's transaction on the pool or the lease: it locks the Session, applies the operation, then prunes the change journal. `BindSession` binds one Session to a transaction's queries as a `SessionTx`, which implements the procedures' interfaces. -Other adapters call only these `sessionpg` participants, on their own transaction's queries: `LockSession`, which orders a write against the Session's Turns; `BindSession`, which runs a Session procedure, such as `CreateEnvironmentDevice`, in that transaction; `LoadEnvironment`, which reads the tenant's Environment of a live Session; and `CreateSessionEnvironment`, which stores the Environment of a Session being created. The other `sessionpg` exports serve `store` until the Session operations leave it. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; `placementpg`, which loads the placement facts and reserves placements, sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. +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` or `TransitionTurn`, in that transaction; `LoadEnvironment`, which reads the tenant's Environment of a live Session; and `CreateSessionEnvironment`, which stores the Environment of a Session being created. The other `sessionpg` exports serve `store` until the Session operations leave it. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; `placementpg`, which loads the placement facts and reserves placements, sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. `store` is transitional. `store.New` builds a pooled Store, and `store.NewExecution(s, lease)` builds the execution writer on a lease it borrows. An execution-only operation on a pooled Store fails with `store.ErrExecutionAuthority`. New adapters do not copy that check: their execution repositories require a `*pgunit.Lease` at construction, their public repositories expose no execution operation, and the check goes away with `store`. @@ -36,7 +36,7 @@ Domain owners, each with its PostgreSQL adapter under `internal/persistence/post - `deployment` and its subpackage `deployment/placement` (`deploymentpg`, `placementpg`): the sandbox deployment and its nodes: provider configuration and the sealed credential, the specification and retained generations, setup, update and switch, node enrollment, identity and authentication, generation configuration, capacity, presence, status, host history, reset, and the counts of nodes and sandboxes bound to the public address; and hosted runtime allocations: reservation with the dedicated device, compute ownership and settlement, observation diagnostics, activity, compute phases, wake receipts and cleanup, with the reads that schedule, discover and authorize them. It interprets Sandbox Provider declarations through the `providers.Registry` it is given, whose lookups return typed errors. `cmd/server` builds that registry and calls it only to build a direct Provider and to discover a Provider's configuration; it takes each setup's mode and declared operations from the `Setup` that `deployment` returns. `deployment/placement` owns hosted admission and placement: `cmd/server` builds one `placement.Rules` from the registry and the public URL, both fixed while Core runs, and gives it to `deployment.Service` and to `store`; its pure decisions admit a hosted Session, choose its node, and admit an allocation's reserved node and a restore on it, and its errors keep one status, code and message through `writeStoreError` and `writeDeploymentError`. `placementpg` loads the facts those decisions read and applies them on the caller's transaction-bound queries; it has no Store or transaction runner. The only other reader is `store`: Session creation decides admission and placement with those rules and `placementpg` inside its own transaction. Deployment changes and allocation writes run through `deployment.ExecutionOperations` on `deploymentpg.NewExecution(lease, …)`, which the Worker receives as `execution.Owner.Deployment`. Each allocation write locks the owning Session, decides in `deployment`, applies in `deploymentpg` and prunes the Session's change journal in the same leased transaction; cleanup settles the Session through the `sessions` procedures on a `sessionpg.SessionTx` bound to that transaction. The Worker reads the deployment, prepares a selection's setup and records live-compute activity through the pooled `deployment.Service` in `execution.Dispatcher.Deployment`, and lists the Sessions a reset still has to archive, schedules node lifecycles and reads allocations through the pooled `deployment.Reader` in `execution.Dispatcher.DeploymentReader`; a pooled read carries no lease, and the leased write that follows it rechecks the allocation's owner. Node management and reads use the pooled `deploymentpg.Store`, which `cmd/server` also reads the owner epoch from and runs the host-history sampler on. Provider calls run outside transactions, and the final transaction rechecks the expected generation. The reset's Session archive stays in `store`. - `coremetrics` (`coremetricspg`): the Core metrics PostgreSQL holds: the root Turn queue counts, the root Turn history, read from one read-only snapshot, and the database size. `cmd/server`'s Core metrics source adds them to the process, pool, Worker and daemon registry measurements. - `runtimehistory` (`runtimehistory/postgresreader`, outside this directory): Runtime history samples: the periodic export, scoped reads and retention, which also prunes node-host samples. -- `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 the archived cancellation receipt the daemon gateway reads), the Sessions whose Runtime the administrator observes across Projects, 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`. +- `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 Turns (their reads and the execution work scan), the model provider a Session froze, which `sessionpg.Store` opens with the credential key, 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, which include a Session's execution binding and the execution device list, their creation, revocation, heartbeats and Runtime enrollment, and the archived cancellation receipt the daemon gateway reads), the Sessions whose Runtime the administrator observes across Projects, 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. Turn transitions, execution completion, which alone publishes or discards a Turn's staged Artifacts, the start of Artifact capture, function calls and their application receipts, the Turn execution journal, Environment initialization, connection observations and their reconciliation, Session device binding, and file-write reservation and settlement run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. ## Request handling diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index ecbb7f82c..6cf3a08ad 100644 --- a/services/core/cmd/server/main.go +++ b/services/core/cmd/server/main.go @@ -406,7 +406,7 @@ func run() error { Sessions: executionStore, SessionCreation: executionStore, SessionEvents: executionStore, - Turns: executionStore, + Turns: sessionStore, Items: sessionStore, Subagents: sessionStore, Artifacts: sessionService, diff --git a/services/core/internal/api/native_classification_integration_test.go b/services/core/internal/api/native_classification_integration_test.go index e9dcde9df..5b0a4530a 100644 --- a/services/core/internal/api/native_classification_integration_test.go +++ b/services/core/internal/api/native_classification_integration_test.go @@ -26,18 +26,14 @@ func TestNativeClassificationPostgresRoundTripAndPublicPrivacy(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err = s.TransitionTurn(t.Context(), tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { - t.Fatal(err) - } + transitionTurn(t, pool, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) status := 503 result := execution.Result{ErrorCode: "engine_failed", Error: "Bearer secret-canary https://private.example/key", EngineErrorCode: code, EngineHTTPStatus: &status, Done: proto.DonePayload{Usage: proto.Usage{InputTokens: 7, OutputTokens: 3}, Metadata: map[string]any{proto.DoneMetaAgentSessionID: "native-secret-canary"}}} outcome, err := json.Marshal(result) if err != nil { t.Fatal(err) } - if _, err = s.TransitionTurn(t.Context(), tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: outcome}); err != nil { - t.Fatal(err) - } + transitionTurn(t, pool, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: outcome}) snap, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil { t.Fatal(err) diff --git a/services/core/internal/api/session_diagnostics_public_compat_test.go b/services/core/internal/api/session_diagnostics_public_compat_test.go index 3f8e5aa7a..6910ec7db 100644 --- a/services/core/internal/api/session_diagnostics_public_compat_test.go +++ b/services/core/internal/api/session_diagnostics_public_compat_test.go @@ -1,6 +1,7 @@ package api import ( + "context" "encoding/json" "net/http" "net/http/httptest" @@ -8,6 +9,7 @@ import ( "strings" "testing" + "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" @@ -66,8 +68,30 @@ func diagnosticRequest(handler http.Handler, path, token string) *httptest.Respo // Item reads from the Session adapter on pool. func databaseSessionReads(s *store.Store, pool *pgxpool.Pool) func(*Dependencies, *testFakes) { return func(d *Dependencies, _ *testFakes) { - d.Sessions, d.Turns, d.SessionAdmin = s, s, s - d.Items = sessionpg.New(pgunit.NewPool(pool), nil) + d.Sessions, d.SessionAdmin = s, s + reads := sessionpg.New(pgunit.NewPool(pool), nil) + d.Items, d.Turns = reads, reads + } +} + +// transitionTurn moves the Turn as the execution owner does, over a pooled +// Session transaction. +func transitionTurn(t *testing.T, pool *pgxpool.Pool, tenant, session, turn string, transition sessions.TurnTransition) { + t.Helper() + tenantID, err := pgunit.ParseID(tenant) + if err != nil { + t.Fatal(err) + } + sessionID, err := pgunit.ParseID(session) + if err != nil { + t.Fatal(err) + } + err = sessionpg.WithSession(t.Context(), pgunit.NewPool(pool), tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { + _, err := sessions.TransitionTurn(ctx, sessionpg.BindSession(q, tenantID, sessionID), turn, transition) + return err + }) + if err != nil { + t.Fatal(err) } } diff --git a/services/core/internal/api/session_diagnostics_test.go b/services/core/internal/api/session_diagnostics_test.go index 572cf8c1a..8fc433ae0 100644 --- a/services/core/internal/api/session_diagnostics_test.go +++ b/services/core/internal/api/session_diagnostics_test.go @@ -23,12 +23,8 @@ func TestDiagnosticsCoreHandlerDatabaseBoundary(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err = s.TransitionTurn(t.Context(), tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { - t.Fatal(err) - } - if _, err = s.TransitionTurn(t.Context(), tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"error_code":"device_disconnected","error":"Bearer raw-secret-canary https://private.example/key","done":{"native_id":"secret-native-canary"}}`)}); err != nil { - t.Fatal(err) - } + transitionTurn(t, pool, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + transitionTurn(t, pool, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"error_code":"device_disconnected","error":"Bearer raw-secret-canary https://private.example/key","done":{"native_id":"secret-native-canary"}}`)}) base := adminSessionsPath + session.ID for _, path := range []string{base + "/diagnostics", base + "/turns/" + receipt.TurnID + "/diagnostics"} { if w := diagnosticRequest(h, path, ""); w.Code != 401 { diff --git a/services/core/internal/execution/archive_cancellation_cleanup_test.go b/services/core/internal/execution/archive_cancellation_cleanup_test.go index 21a4a3916..eda95d545 100644 --- a/services/core/internal/execution/archive_cancellation_cleanup_test.go +++ b/services/core/internal/execution/archive_cancellation_cleanup_test.go @@ -99,7 +99,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { // This fixture isolates lifecycle ordering. Protocol-driven waiting is // independently exercised in TestArchiveWaitingCancellationReceipts. for _, transition := range []sessions.TurnTransition{{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}, {ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnWaiting}} { - if _, err := writer.TransitionTurn(t.Context(), project.TenantID, session.ID, input.TurnID, transition); err != nil { + if _, err := leased.Sessions.TransitionTurn(t.Context(), project.TenantID, session.ID, input.TurnID, transition); err != nil { t.Fatal(err) } } @@ -155,11 +155,12 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { if err != nil { t.Fatal(err) } + sessionReader, _ := testSessions(t, pool, testCredentialCipher(t)) kills := 0 expectedStatus := sessions.TurnWaiting provider := waitingCleanupProvider{beforeKill: func() { kills++ - turn, err := s.GetTurn(t.Context(), project.TenantID, session.ID, input.TurnID) + turn, err := sessionReader.GetTurn(t.Context(), project.TenantID, session.ID, input.TurnID) if err != nil || turn.Status != expectedStatus || turn.CancelRequestedAt.IsZero() || (turn.CompletedAt.IsZero() != (expectedStatus == sessions.TurnWaiting)) { t.Fatal("cleanup observed unexpected terminal state", turn, err) } @@ -168,7 +169,6 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { t.Fatal("Kill bypassed durable cleanup ownership", allocation, err) } }} - sessionReader, _ := testSessions(t, pool, testCredentialCipher(t)) lifecycle := &runtimeLifecycle{store: writer, sessions: sessionReader, sessionExecution: leased.Sessions, deployment: leased.Deployment, deployments: deployments, reader: reader, lease: leased.Lease, registry: registry, config: RuntimeProvider{InstallationID: installation, Provider: provider}, connections: map[string]*runtimeConnection{}} if checkpoint { lifecycle.config.Provider = waitingCleanupCheckpoint{beforeKill: provider.beforeKill} @@ -185,7 +185,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { t.Fatal(pending, err) } // Controlled terminal receipt fixture; no native cancellation claim. - if _, err := writer.TransitionTurn(t.Context(), project.TenantID, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnWaiting, Status: sessions.TurnCancelled}); err != nil { + if _, err := leased.Sessions.TransitionTurn(t.Context(), project.TenantID, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnWaiting, Status: sessions.TurnCancelled}); err != nil { t.Fatal(err) } expectedStatus = sessions.TurnCancelled @@ -197,7 +197,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { if err != nil || after.State != "released" || kills != 1 { t.Fatal("cleanup did not release", after, kills, err) } - turn, err := s.GetTurn(t.Context(), project.TenantID, session.ID, input.TurnID) + turn, err := sessionReader.GetTurn(t.Context(), project.TenantID, session.ID, input.TurnID) if err != nil || turn.Status != expectedStatus || (turn.CompletedAt.IsZero() != (expectedStatus == sessions.TurnWaiting)) { t.Fatal("cleanup should not fabricate cancellation", turn, err) } diff --git a/services/core/internal/execution/artifacts.go b/services/core/internal/execution/artifacts.go index 203f5ada1..f1aa98504 100644 --- a/services/core/internal/execution/artifacts.go +++ b/services/core/internal/execution/artifacts.go @@ -17,7 +17,7 @@ func (d *Dispatcher) captureCompletedArtifacts(ctx context.Context, peer *runtim } owner, cancel := context.WithTimeout(ctx, 180*time.Second) defer cancel() - err := d.Store.BeginTurnArtifactCapture(owner, session.TenantID, session.ID, turnID, result.AppliedThrough) + err := d.sessionExecution.BeginTurnArtifactCapture(owner, session.TenantID, session.ID, turnID, result.AppliedThrough) if err == nil { err = d.withPreparedWorkspace(owner, peer, session, environment, bound, func(ctx context.Context, handle string) error { return peer.ExportWorkspaceOutputs(ctx, proto.WorkspaceExportPayload{Handle: handle, EnvironmentID: environment.ID}, func(body io.Reader) error { @@ -36,7 +36,7 @@ func (d *Dispatcher) captureCompletedArtifacts(ctx context.Context, peer *runtim status = sessions.TurnFailed check, stop := context.WithTimeout(context.Background(), 5*time.Second) defer stop() - if turn, err := d.Store.GetTurn(check, session.TenantID, session.ID, turnID); err == nil && !turn.CancelRequestedAt.IsZero() { + if turn, err := d.SessionsReader.GetTurn(check, session.TenantID, session.ID, turnID); err == nil && !turn.CancelRequestedAt.IsZero() { status = sessions.TurnCancelled } return result, status diff --git a/services/core/internal/execution/delivery.go b/services/core/internal/execution/delivery.go index a9c2ea6d6..ab20d7d7c 100644 --- a/services/core/internal/execution/delivery.go +++ b/services/core/internal/execution/delivery.go @@ -105,7 +105,7 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe var pending *pendingInput var cancelSent time.Time var cancelReply <-chan cancellationResult - functions := &functionExchange{kind: request.AgentKind, store: d.Store, sessions: d.sessionExecution, tenant: tenantID, session: sessionID, turn: request.RunID, tools: request.FunctionTools} + functions := &functionExchange{kind: request.AgentKind, turns: d.SessionsReader, sessions: d.sessionExecution, tenant: tenantID, session: sessionID, turn: request.RunID, tools: request.FunctionTools} done := false cancelCtx, stopCancellation := context.WithCancel(ctx) defer stopCancellation() @@ -278,7 +278,7 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe } continue } - turn, err := d.Store.GetTurn(ctx, tenantID, sessionID, request.RunID) + turn, err := d.SessionsReader.GetTurn(ctx, tenantID, sessionID, request.RunID) if err != nil { result.ErrorCode = "execution_state_unavailable" if ctx.Err() != nil { diff --git a/services/core/internal/execution/deployment_provider_observations_test.go b/services/core/internal/execution/deployment_provider_observations_test.go index 2816644ea..fa11d22b2 100644 --- a/services/core/internal/execution/deployment_provider_observations_test.go +++ b/services/core/internal/execution/deployment_provider_observations_test.go @@ -16,6 +16,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/modelconfiguration" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/modelconfigurationpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" @@ -24,7 +25,7 @@ import ( type finishObservationFixture struct { s *store.Store - writer *store.Store + execution *sessions.ExecutionOperations lease Ownership pool *pgxpool.Pool defaults *modelconfigurationpg.Store @@ -77,7 +78,7 @@ func newFinishObservationFixture(t *testing.T, maxConnections int32) finishObser if err != nil { t.Fatal(err) } - return finishObservationFixture{s, owner.Store, owner.Lease, pool, defaults, tenant, session, *dispatcher} + return finishObservationFixture{s, owner.Sessions, owner.Lease, pool, defaults, tenant, session, *dispatcher} } func (f finishObservationFixture) start(t *testing.T) sessions.InputReceipt { t.Helper() @@ -85,7 +86,7 @@ func (f finishObservationFixture) start(t *testing.T) sessions.InputReceipt { if err != nil { t.Fatal(err) } - if _, err = f.writer.TransitionTurn(t.Context(), f.tenant, f.session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err = f.execution.TransitionTurn(t.Context(), f.tenant, f.session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } return receipt @@ -194,7 +195,7 @@ func TestFinishRunObservationLockTimeoutAndFailureKeepLease(t *testing.T) { t.Fatal(err) } defer conn.Release() - turn, err := f.writer.CompleteExecution(t.Context(), f.tenant, f.session.ID, receipt.TurnID, sessions.TurnCompleted, json.RawMessage(`{}`), "", receipt.Sequence) + turn, err := f.execution.CompleteExecution(t.Context(), f.tenant, f.session.ID, receipt.TurnID, sessions.TurnCompleted, json.RawMessage(`{}`), "", receipt.Sequence) if err != nil { t.Fatal(err) } @@ -224,7 +225,7 @@ func TestFinishRunObservationLockTimeoutAndFailureKeepLease(t *testing.T) { if cleanup != nil { cleanup() } - persisted, err := f.s.GetTurn(t.Context(), f.tenant, f.session.ID, receipt.TurnID) + persisted, err := sessionpg.New(pgunit.NewPool(f.pool), nil).GetTurn(t.Context(), f.tenant, f.session.ID, receipt.TurnID) if err != nil || persisted.Status != sessions.TurnCompleted { t.Fatal("terminal outcome lost", err) } diff --git a/services/core/internal/execution/dispatcher.go b/services/core/internal/execution/dispatcher.go index 54ad6433d..12710583e 100644 --- a/services/core/internal/execution/dispatcher.go +++ b/services/core/internal/execution/dispatcher.go @@ -79,7 +79,7 @@ func (d *Dispatcher) Run(ctx context.Context, tenantID, sessionID, turnID string if err != nil { return sessions.Turn{}, err } - bound, err := d.Store.GetSessionExecutionBinding(ctx, tenantID, sessionID) + bound, err := d.SessionsReader.GetSessionExecutionBinding(ctx, tenantID, sessionID) if err != nil { return sessions.Turn{}, err } @@ -115,7 +115,7 @@ func (d *Dispatcher) Run(ctx context.Context, tenantID, sessionID, turnID string return sessions.Turn{}, err } defer prepared.close() - if _, err := d.Store.TransitionTurn(ctx, tenantID, sessionID, turnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err := d.sessionExecution.TransitionTurn(ctx, tenantID, sessionID, turnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { return sessions.Turn{}, err } req.ConversationID, req.RunID, req.Input = sessionID, turnID, text @@ -141,11 +141,11 @@ func (d *Dispatcher) finishRun(tenantID, sessionID, turnID, model string, result encoded, _ = json.Marshal(result) status, nativeID = sessions.TurnFailed, "" } - turn, err := d.Store.CompleteExecution(finishCtx, tenantID, sessionID, turnID, status, encoded, nativeID, result.AppliedThrough) + turn, err := d.sessionExecution.CompleteExecution(finishCtx, tenantID, sessionID, turnID, status, encoded, nativeID, result.AppliedThrough) if errors.Is(err, sessions.ErrUnappliedInputs) { result.ErrorCode = "input_not_applied" encoded, _ = json.Marshal(result) - turn, err = d.Store.CompleteExecution(finishCtx, tenantID, sessionID, turnID, sessions.TurnFailed, encoded, nativeID, result.AppliedThrough) + turn, err = d.sessionExecution.CompleteExecution(finishCtx, tenantID, sessionID, turnID, sessions.TurnFailed, encoded, nativeID, result.AppliedThrough) } if err == nil { d.observeDeploymentProvider(tenantID, sessionID, turn) diff --git a/services/core/internal/execution/executor_preparation.go b/services/core/internal/execution/executor_preparation.go index fad7e9c36..11fafdf9d 100644 --- a/services/core/internal/execution/executor_preparation.go +++ b/services/core/internal/execution/executor_preparation.go @@ -37,7 +37,7 @@ func (d *Dispatcher) awaitTurnExecutor(ctx context.Context, tenant, session, tur case <-ctx.Done(): return ctx.Err() case <-tick.C: - current, err := d.Store.GetTurn(ctx, tenant, session, turn) + current, err := d.SessionsReader.GetTurn(ctx, tenant, session, turn) if err != nil { return err } diff --git a/services/core/internal/execution/functions.go b/services/core/internal/execution/functions.go index ec278de1a..f3804de11 100644 --- a/services/core/internal/execution/functions.go +++ b/services/core/internal/execution/functions.go @@ -12,7 +12,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/items" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func functionTools(raw []json.RawMessage) ([]proto.FunctionTool, error) { @@ -50,7 +49,7 @@ type functionReply struct { } type functionExchange struct { - store *store.Store + turns sessions.TurnReader sessions *sessions.ExecutionOperations tenant, session, turn string kind string @@ -125,7 +124,7 @@ func (f *functionExchange) confirm(ctx context.Context, reply functionReply) err func (f *functionExchange) unlessCancelling(ctx context.Context, err error) error { if errors.Is(err, sessions.ErrTurnConflict) { - turn, lookupErr := f.store.GetTurn(ctx, f.tenant, f.session, f.turn) + turn, lookupErr := f.turns.GetTurn(ctx, f.tenant, f.session, f.turn) if lookupErr == nil && !turn.CancelRequestedAt.IsZero() { return nil } @@ -144,7 +143,7 @@ func (f *functionExchange) complete(ctx context.Context) error { if len(calls) != 0 { return errors.New("function calls remain unapplied") } - turn, err := f.store.GetTurn(ctx, f.tenant, f.session, f.turn) + turn, err := f.turns.GetTurn(ctx, f.tenant, f.session, f.turn) if err != nil { return err } diff --git a/services/core/internal/execution/model_execution.go b/services/core/internal/execution/model_execution.go index ad02143c7..926aac7e5 100644 --- a/services/core/internal/execution/model_execution.go +++ b/services/core/internal/execution/model_execution.go @@ -2,17 +2,13 @@ package execution import ( "context" - "errors" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) func (d *Dispatcher) sessionModelOptions(ctx context.Context, session sessions.Session) (map[string]any, error) { - if d.Store == nil { - return nil, errors.New("session model configuration is unavailable") - } - provider, err := d.Store.SessionModelExecution(ctx, session.TenantID, session.ID) + provider, err := d.SessionsReader.SessionModelExecution(ctx, session.TenantID, session.ID) if err != nil { return nil, err } diff --git a/services/core/internal/execution/model_execution_test.go b/services/core/internal/execution/model_execution_test.go index f6fd32bc6..627a3161f 100644 --- a/services/core/internal/execution/model_execution_test.go +++ b/services/core/internal/execution/model_execution_test.go @@ -6,15 +6,20 @@ import ( "strings" "testing" + "github.com/google/uuid" + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) func TestSessionModelExecutionNeverFallsBack(t *testing.T) { - var d Dispatcher - if _, err := d.executionRequest(t.Context(), sessions.Session{Engine: "codex"}, Snapshot{ModelProviderConfigured: true}, runtimedevice.KindCapabilities{}, sessions.ExecutionBinding{}); err == nil { - t.Fatal("missing Session credentials fell back") + reader, _ := testSessions(t, pgtest.Open(t), nil) + d := Dispatcher{SessionsReader: reader} + session := sessions.Session{TenantID: uuid.NewString(), ID: uuid.NewString(), Engine: "codex"} + if _, err := d.executionRequest(t.Context(), session, Snapshot{ModelProviderConfigured: true}, runtimedevice.KindCapabilities{}, sessions.ExecutionBinding{}); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal("missing Session credentials fell back", err) } // Hosted and self-hosted Runtimes have no model configuration of their own. for _, environment := range []string{"openai_hosted", "self_hosted"} { diff --git a/services/core/internal/execution/owner_test.go b/services/core/internal/execution/owner_test.go index 2f5d3eae2..d62f6e18f 100644 --- a/services/core/internal/execution/owner_test.go +++ b/services/core/internal/execution/owner_test.go @@ -53,8 +53,8 @@ func (l *closeCountingLease) Close(ctx context.Context) error { return l.inner.Close(ctx) } -// unusedSessions returns the Session service and reader for Workers that run -// no Turn. Any call through them panics. +// unusedSessions returns the Session service and reader for Workers that fail +// to start. Any call through them panics. func unusedSessions(t *testing.T) (*sessions.Service, sessions.Reader) { t.Helper() service, err := sessions.NewService(struct{ sessions.Storage }{}) @@ -195,10 +195,11 @@ func TestStartWorkerChecksDeploymentAfterItsDependencies(t *testing.T) { } func TestWorkerRunClosesLeaseAfterDrain(t *testing.T) { - s, owner, deployments, deploymentReader := resetManagerStore(t) + s, owner, deployments, deploymentReader, pool := resetManagerStoreDB(t, nil) lease := &closeCountingLease{t: t, inner: owner.Lease} id := uuid.NewString() - service, reader := unusedSessions(t) + // The Worker's first reconciliation scans the Session work. + reader, service := testSessions(t, pool, nil) dispatcher := &Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), Credentials: &recordingCredentials{}, Observer: unusedObserver{t}, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader, ManagedRuntimes: NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })} worker, err := StartWorker(t.Context(), dispatcher, Owner{Lease: lease, Store: owner.Store, Deployment: owner.Deployment, Sessions: owner.Sessions}) if err != nil { diff --git a/services/core/internal/execution/prepared_dispatch.go b/services/core/internal/execution/prepared_dispatch.go index 204a3d11b..c9dec7ede 100644 --- a/services/core/internal/execution/prepared_dispatch.go +++ b/services/core/internal/execution/prepared_dispatch.go @@ -42,7 +42,7 @@ func (d *Dispatcher) RunEnvironmentInput(ctx context.Context, lease Ownership, t // Reserved before providers were required; the caller settles it as failed. return run, ErrModelProviderRequired } - bound, err := d.Store.GetSessionExecutionBinding(ctx, tenantID, sessionID) + bound, err := d.SessionsReader.GetSessionExecutionBinding(ctx, tenantID, sessionID) if err != nil { return run, err } diff --git a/services/core/internal/execution/worker.go b/services/core/internal/execution/worker.go index 2f584c543..887722017 100644 --- a/services/core/internal/execution/worker.go +++ b/services/core/internal/execution/worker.go @@ -352,7 +352,7 @@ func (w *Worker) Run(ctx context.Context) (runErr error) { func (w *Worker) reconcile(ctx context.Context) error { cursor := "" for { - work, err := w.dispatcher.Store.ListExecutionWork(ctx, cursor, []string{sessions.TurnInProgress, sessions.TurnWaiting}, nil) + work, err := w.dispatcher.SessionsReader.ListExecutionWork(ctx, cursor, []string{sessions.TurnInProgress, sessions.TurnWaiting}, nil) if err != nil { return err } @@ -360,7 +360,7 @@ func (w *Worker) reconcile(ctx context.Context) error { return nil } for _, item := range work { - _, err := w.dispatcher.Store.TransitionTurn(ctx, item.TenantID, item.SessionID, item.TurnID, sessions.TurnTransition{ExpectedStatus: item.Status, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"error_code":"execution_interrupted"}`)}) + _, err := w.dispatcher.sessionExecution.TransitionTurn(ctx, item.TenantID, item.SessionID, item.TurnID, sessions.TurnTransition{ExpectedStatus: item.Status, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"error_code":"execution_interrupted"}`)}) if err != nil && !errors.Is(err, sessions.ErrTurnConflict) { return err } @@ -382,7 +382,7 @@ func (w *Worker) runClaim(ctx context.Context, item sessions.ExecutionWork) erro } finish, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - turn, err := w.dispatcher.Store.GetTurn(finish, item.TenantID, item.SessionID, item.TurnID) + turn, err := w.dispatcher.SessionsReader.GetTurn(finish, item.TenantID, item.SessionID, item.TurnID) if err != nil { return err } @@ -395,7 +395,7 @@ func (w *Worker) runClaim(ctx context.Context, item sessions.ExecutionWork) erro return nil } log.Ctx(ctx).Error("oac-core dispatch did not complete", "turn_id", item.TurnID) - _, err = w.dispatcher.Store.TransitionTurn(finish, item.TenantID, item.SessionID, item.TurnID, sessions.TurnTransition{ExpectedStatus: turn.Status, Status: sessions.TurnFailed, Outcome: outcome}) + _, err = w.dispatcher.sessionExecution.TransitionTurn(finish, item.TenantID, item.SessionID, item.TurnID, sessions.TurnTransition{ExpectedStatus: turn.Status, Status: sessions.TurnFailed, Outcome: outcome}) if errors.Is(err, sessions.ErrTurnConflict) { return nil } diff --git a/services/core/internal/execution/worker_device.go b/services/core/internal/execution/worker_device.go index 9e7cc9e16..ea63a6b9f 100644 --- a/services/core/internal/execution/worker_device.go +++ b/services/core/internal/execution/worker_device.go @@ -18,7 +18,7 @@ func (w *Worker) bind(ctx context.Context, item sessions.ExecutionWork) (bool, e } // Candidate selection is a snapshot. Cancellation can append a control input // before this read, so recheck eligibility after reading the input history. - turn, err := w.dispatcher.Store.GetTurn(ctx, item.TenantID, item.SessionID, item.TurnID) + turn, err := w.dispatcher.SessionsReader.GetTurn(ctx, item.TenantID, item.SessionID, item.TurnID) if errors.Is(err, sessions.ErrNotFound) { return false, nil } @@ -38,7 +38,7 @@ func (w *Worker) bind(ctx context.Context, item sessions.ExecutionWork) (bool, e if !errors.Is(err, sessions.ErrDeviceBindingConflict) { return ready, err } - _, err = w.dispatcher.Store.TransitionTurn(ctx, item.TenantID, item.SessionID, item.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"error_code":"execution_device_unavailable"}`)}) + _, err = w.dispatcher.sessionExecution.TransitionTurn(ctx, item.TenantID, item.SessionID, item.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"error_code":"execution_device_unavailable"}`)}) if errors.Is(err, sessions.ErrTurnConflict) { err = nil } @@ -107,7 +107,7 @@ func (w *Worker) bindSessionDevice(ctx context.Context, session sessions.Session if !errors.Is(err, sessions.ErrNotFound) { return false, err } - devices, err := w.dispatcher.Store.ListExecutionDevices(ctx, session.TenantID) + devices, err := w.dispatcher.SessionsReader.ListExecutionDevices(ctx, session.TenantID) if err != nil { return false, err } diff --git a/services/core/internal/execution/worker_schedule.go b/services/core/internal/execution/worker_schedule.go index 2a901f5f2..fbc0098d8 100644 --- a/services/core/internal/execution/worker_schedule.go +++ b/services/core/internal/execution/worker_schedule.go @@ -21,13 +21,13 @@ type scheduledWork struct { } func (s *workerSchedule) selectWork(ctx context.Context, w *Worker, devices []string, active map[string]bool) ([]scheduledWork, error) { - turns, err := w.dispatcher.Store.ListExecutionWork(ctx, s.turnCursor, []string{sessions.TurnQueued}, devices) + turns, err := w.dispatcher.SessionsReader.ListExecutionWork(ctx, s.turnCursor, []string{sessions.TurnQueued}, devices) if err != nil { return nil, err } if len(turns) == 0 && s.turnCursor != "" { s.turnCursor = "" - turns, err = w.dispatcher.Store.ListExecutionWork(ctx, "", []string{sessions.TurnQueued}, devices) + turns, err = w.dispatcher.SessionsReader.ListExecutionWork(ctx, "", []string{sessions.TurnQueued}, devices) if err != nil { return nil, err } diff --git a/services/core/internal/persistence/postgres/sessionpg/binding.go b/services/core/internal/persistence/postgres/sessionpg/binding.go index 04fadf5bd..cf227e3f5 100644 --- a/services/core/internal/persistence/postgres/sessionpg/binding.go +++ b/services/core/internal/persistence/postgres/sessionpg/binding.go @@ -34,9 +34,8 @@ var ( _ sessions.InputStartTx = (*SessionTx)(nil) _ sessions.ComputeAdmissionTx = (*SessionTx)(nil) _ sessions.EnvironmentDeviceTx = (*SessionTx)(nil) - _ sessions.TurnJournalTx = (*SessionTx)(nil) - _ sessions.TurnEventTx = (*SessionTx)(nil) _ sessions.InputProjectionTx = (*SessionTx)(nil) + _ sessions.TurnTx = (*SessionTx)(nil) ) // BindSession binds the tenant's Session to the caller's transaction-bound @@ -70,6 +69,9 @@ func (t *SessionTx) LoadTurn(ctx context.Context, turn string) (sessions.Turn, e return sessions.Turn{}, err } row, err := t.q.SessionEventTurn(ctx, sqlc.SessionEventTurnParams{SessionID: t.session, ID: id}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.Turn{}, sessions.ErrNotFound + } if err != nil { return sessions.Turn{}, err } diff --git a/services/core/internal/persistence/postgres/sessionpg/devices.go b/services/core/internal/persistence/postgres/sessionpg/devices.go index 0db12ccb9..9bf5e67e5 100644 --- a/services/core/internal/persistence/postgres/sessionpg/devices.go +++ b/services/core/internal/persistence/postgres/sessionpg/devices.go @@ -20,14 +20,60 @@ func (s *Store) GetSessionDevice(ctx context.Context, tenant, session string) (s if err != nil { return sessions.ExecutionDevice{}, err } - ready, err := s.units.Queries().GetSessionInitializationReady(ctx, sqlc.GetSessionInitializationReadyParams(lookup)) + if err := requireInitialized(ctx, s.units.Queries(), lookup); err != nil { + return sessions.ExecutionDevice{}, err + } + return s.GetSessionRuntimeDevice(ctx, tenant, session) +} + +func (s *Store) GetSessionExecutionBinding(ctx context.Context, tenant, session string) (sessions.ExecutionBinding, error) { + lookup, err := DeviceLookup(tenant, session) + if err != nil { + return sessions.ExecutionBinding{}, err + } + q := s.units.Queries() + if err := requireInitialized(ctx, q, lookup); err != nil { + return sessions.ExecutionBinding{}, err + } + row, err := q.GetSessionExecutionBinding(ctx, sqlc.GetSessionExecutionBindingParams(lookup)) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.ExecutionBinding{}, sessions.ErrNotFound + } + if err != nil { + return sessions.ExecutionBinding{}, err + } + return sessions.ExecutionBinding{ + Device: sessions.ExecutionDevice{ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name, EnvironmentID: optionalID(row.EnvironmentID)}, + NativeSessionID: row.NativeSessionID, + HasStartedTurn: row.HasStartedTurn, + }, nil +} + +// requireInitialized requires that the tenant's Session completed its +// Environment preparation; before that, and for a missing Session, it is +// sessions.ErrNotFound. +func requireInitialized(ctx context.Context, q *sqlc.Queries, lookup sqlc.GetDeviceParams) error { + ready, err := q.GetSessionInitializationReady(ctx, sqlc.GetSessionInitializationReadyParams(lookup)) if errors.Is(err, pgx.ErrNoRows) || (err == nil && !ready) { - return sessions.ExecutionDevice{}, sessions.ErrNotFound + return sessions.ErrNotFound } + return err +} + +func (s *Store) ListExecutionDevices(ctx context.Context, tenant string) ([]sessions.ExecutionDevice, error) { + id, err := parseID(tenant) if err != nil { - return sessions.ExecutionDevice{}, err + return nil, err } - return s.GetSessionRuntimeDevice(ctx, tenant, session) + rows, err := s.units.Queries().ListExecutionDevices(ctx, id) + if err != nil { + return nil, err + } + devices := make([]sessions.ExecutionDevice, 0, len(rows)) + for _, row := range rows { + devices = append(devices, sessions.ExecutionDevice{ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name}) + } + return devices, nil } func (s *Store) GetSessionRuntimeDevice(ctx context.Context, tenant, session string) (sessions.ExecutionDevice, error) { diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_journal.go b/services/core/internal/persistence/postgres/sessionpg/execution_journal.go deleted file mode 100644 index ad0f9b1a1..000000000 --- a/services/core/internal/persistence/postgres/sessionpg/execution_journal.go +++ /dev/null @@ -1,24 +0,0 @@ -package sessionpg - -import ( - "context" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -// WithTurnJournal runs apply in a Session transaction on the lease, including -// for a publicly deleted Session. -func (e *Execution) WithTurnJournal(ctx context.Context, tenantID, sessionID string, apply func(context.Context, sessions.TurnJournalTx) error) error { - tenant, err := parseID(tenantID) - if err != nil { - return err - } - session, err := parseID(sessionID) - if err != nil { - return err - } - return WithSession(ctx, e.lease, tenant, session, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { - return apply(ctx, BindSession(q, tenant, session)) - }) -} diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_turns.go b/services/core/internal/persistence/postgres/sessionpg/execution_turns.go new file mode 100644 index 000000000..6afe1e0bc --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution_turns.go @@ -0,0 +1,80 @@ +package sessionpg + +import ( + "context" + "errors" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// WithTurns runs apply in a Session transaction on the lease, including for a +// publicly deleted Session. +func (e *Execution) WithTurns(ctx context.Context, tenantID, sessionID string, apply func(context.Context, sessions.TurnTx) error) error { + tenant, err := parseID(tenantID) + if err != nil { + return err + } + session, err := parseID(sessionID) + if err != nil { + return err + } + return WithSession(ctx, e.lease, tenant, session, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { + return apply(ctx, BindSession(q, tenant, session)) + }) +} + +func (t *SessionTx) ApplyTurnStatus(ctx context.Context, turn string, change sessions.TurnStatusChange) (sessions.Turn, error) { + id, err := parseID(turn) + if err != nil { + return sessions.Turn{}, err + } + row, err := t.q.TransitionTurn(ctx, sqlc.TransitionTurnParams{ + ID: id, SessionID: t.session, ExpectedStatus: change.Expected, NewStatus: change.Status, Outcome: change.Outcome, + SourceCompletedAt: pgtype.Timestamptz{Time: change.SourceCompletedAt, Valid: !change.SourceCompletedAt.IsZero()}, + }) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.Turn{}, sessions.ErrTurnConflict + } + if err != nil { + return sessions.Turn{}, err + } + return TurnFromRow(row), nil +} + +func (t *SessionTx) HasUnappliedInputs(ctx context.Context, turn string, appliedThrough int64) (bool, error) { + id, err := parseID(turn) + if err != nil { + return false, err + } + return t.q.HasUnappliedMessages(ctx, sqlc.HasUnappliedMessagesParams{SessionID: t.session, TurnID: id, Sequence: appliedThrough}) +} + +func (t *SessionTx) RememberNativeSession(ctx context.Context, native string) error { + n, err := t.q.RememberNativeSession(ctx, sqlc.RememberNativeSessionParams{SessionID: t.session, NativeSessionID: native}) + if err != nil { + return err + } + if n != 1 { + return sessions.ErrNotFound + } + return nil +} + +func (t *SessionTx) BeginArtifactCapture(ctx context.Context, turn string) error { + id, err := parseID(turn) + if err != nil { + return err + } + n, err := t.q.BeginTurnArtifactCapture(ctx, sqlc.BeginTurnArtifactCaptureParams{SessionID: t.session, ID: id}) + if err != nil { + return err + } + if n != 1 { + return sessions.ErrTurnConflict + } + return nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_turns_test.go b/services/core/internal/persistence/postgres/sessionpg/execution_turns_test.go new file mode 100644 index 000000000..6e06a2780 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution_turns_test.go @@ -0,0 +1,206 @@ +package sessionpg + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "reflect" + "slices" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// transition moves the Turn as the seeder does: the TransitionTurn procedure +// in a pooled Session transaction. +func transition(ctx context.Context, pool *pgxpool.Pool, tenant, session pgtype.UUID, turn string, transition sessions.TurnTransition) (sessions.Turn, error) { + var moved sessions.Turn + err := WithSession(ctx, pgunit.NewPool(pool), tenant, session, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { + var err error + moved, err = sessions.TransitionTurn(ctx, BindSession(q, tenant, session), turn, transition) + return err + }) + return moved, err +} + +// Concurrent terminal callbacks commit exactly one outcome, which a late +// callback and a new pool see unchanged. +func TestConcurrentTerminalTransitionsKeepOneOutcome(t *testing.T) { + pool := pgtest.Open(t) + tenant, session, id := newTurn(t, pool, sessions.TurnInProgress) + turn := text(id) + exec(t, pool, `UPDATE turns SET started_at = clock_timestamp(), cancel_requested_at = clock_timestamp() WHERE id = $1`, id) + var wg sync.WaitGroup + winners, errs := make(chan sessions.Turn, 3), make(chan error, 3) + for _, status := range []string{sessions.TurnCompleted, sessions.TurnFailed, sessions.TurnCancelled} { + wg.Go(func() { + outcome := json.RawMessage(`{"reported":"` + status + `"}`) + moved, err := transition(t.Context(), pool, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: status, Outcome: outcome}) + if err != nil { + errs <- err + return + } + winners <- moved + }) + } + wg.Wait() + close(winners) + close(errs) + if len(winners) != 1 || len(errs) != 2 { + t.Fatalf("winners=%d errors=%d", len(winners), len(errs)) + } + for err := range errs { + if !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatal(err) + } + } + winner := <-winners + if winner.CompletedAt.IsZero() || winner.CancelRequestedAt.IsZero() || winner.CompletedAt.Before(winner.StartedAt) { + t.Fatalf("terminal timestamps: %+v", winner) + } + if _, err := transition(t.Context(), pool, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"late":true}`)}); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatalf("late terminal callback accepted: %v", err) + } + got, err := New(pgunit.NewPool(pgtest.Open(t)), nil).GetTurn(t.Context(), text(tenant), text(session), turn) + if err != nil || !reflect.DeepEqual(got, winner) { + t.Fatalf("terminal outcome changed: %+v, %v", got, err) + } +} + +// Artifact capture and completion admit only applied inputs. A completion that +// fails rolls back whole; one that commits records its event, publishes the +// staged Artifacts and journals the Turn's end in order. A closed lease +// writes nothing. +func TestCompleteExecutionSettlesTheTurnOnTheLease(t *testing.T) { + pool := pgtest.OpenIsolated(t, nil) + operations, lease := sessionExecution(t, pool) + _, service := stagingService(t, pool) + tenant, session, environment, turn := stagingTurn(t, pool, "openai_hosted") + if err := stage(t, service, tenant, session, turn, environment, bytes.NewReader(artifactExport(t, [2]string{"outputs/a", "data"}))); err != nil { + t.Fatal(err) + } + var applied int64 + if err := pool.QueryRow(t.Context(), `INSERT INTO turn_inputs(session_id, turn_id, idempotency_key, kind, payload) + VALUES ($1, $2, 'message', 'message', '{}') RETURNING sequence`, session, turn).Scan(&applied); err != nil { + t.Fatal(err) + } + if err := operations.BeginTurnArtifactCapture(t.Context(), tenant, session, turn, applied-1); !errors.Is(err, sessions.ErrUnappliedInputs) { + t.Fatal("capture before the inputs applied", err) + } + if err := operations.BeginTurnArtifactCapture(t.Context(), tenant, session, turn, applied); err != nil { + t.Fatal(err) + } + if err := operations.BeginTurnArtifactCapture(t.Context(), tenant, session, turn, applied); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatal("second capture", err) + } + + completedAt := time.Now().Add(-time.Hour).Truncate(time.Millisecond) + outcome := json.RawMessage(fmt.Sprintf(`{"done":{"source_completed_at_ms":%d}}`, completedAt.UnixMilli())) + _, before := journal(t, pool, pgunit.PathID(session)) + if _, err := operations.CompleteExecution(t.Context(), tenant, session, turn, sessions.TurnCompleted, outcome, "", applied-1); !errors.Is(err, sessions.ErrUnappliedInputs) { + t.Fatal("completion before the inputs applied", err) + } + // The Session has no bound device to remember the native session on. + if _, err := operations.CompleteExecution(t.Context(), tenant, session, turn, sessions.TurnCompleted, outcome, "native", applied); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal("completion without a bound device", err) + } + if entries, count, _ := recorded(t, pool, pgunit.PathID(turn)); len(entries) != 0 || count != 0 { + t.Fatalf("rolled back completion kept %q", entries) + } + if _, after := journal(t, pool, pgunit.PathID(session)); len(after) != len(before) { + t.Fatalf("rolled back completion journaled %q", kinds(after[len(before):])) + } + + ended, err := operations.CompleteExecution(t.Context(), tenant, session, turn, sessions.TurnCompleted, outcome, "", applied) + if err != nil || ended.Status != sessions.TurnCompleted || !ended.CompletedAt.Equal(completedAt) { + t.Fatalf("completion: %+v, %v", ended, err) + } + if entries, _, _ := recorded(t, pool, pgunit.PathID(turn)); !reflect.DeepEqual(entries, []string{"execution_completed"}) { + t.Fatalf("turn journal %q", entries) + } + var published int + if err := pool.QueryRow(t.Context(), `SELECT count(*) FROM session_artifacts WHERE turn_id = $1 AND created_at IS NOT NULL`, turn).Scan(&published); err != nil || published != 1 { + t.Fatalf("published %d Artifacts: %v", published, err) + } + _, after := journal(t, pool, pgunit.PathID(session)) + if got, want := kinds(after[len(before):]), []string{"agent.session.turn.completed", "agent.session.idle"}; !reflect.DeepEqual(got, want) { + t.Fatalf("session journal %q, want %q", got, want) + } + + next := text(addTurn(t, pool, pgunit.PathID(session), sessions.TurnQueued)) + if err := lease.Close(t.Context()); err != nil { + t.Fatal(err) + } + if _, err := operations.TransitionTurn(t.Context(), tenant, session, next, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnCancelled}); !errors.Is(err, pgunit.ErrLeaseClosed) { + t.Fatal("closed lease moved a Turn", err) + } + if _, err := operations.CompleteExecution(t.Context(), tenant, session, next, sessions.TurnFailed, nil, "", 0); !errors.Is(err, pgunit.ErrLeaseClosed) { + t.Fatal("closed lease completed a Turn", err) + } +} + +// Turn pages keep creation order with an ID tie-breaker in both directions, +// and a cursor or Session outside the tenant's Session is missing. +func TestTurnPagesKeepScopeAndOrder(t *testing.T) { + pool := pgtest.Open(t) + store := New(pgunit.NewPool(pool), nil) + ctx := t.Context() + tenantID, sessionID, first := newTurn(t, pool, sessions.TurnCancelled) + tenant, session := text(tenantID), text(sessionID) + ids := []string{text(first)} + for range 3 { + ids = append(ids, text(addTurn(t, pool, sessionID, sessions.TurnCancelled))) + } + // Equal creation times exercise the ID tie-breaker across page boundaries. + exec(t, pool, `UPDATE turns SET created_at = $1 WHERE session_id = $2`, time.Unix(1700000000, 0), sessionID) + slices.Sort(ids) + for _, ascending := range []bool{true, false} { + cursor := "" + for i := range ids { + page, err := store.ListTurns(ctx, tenant, session, cursor, 1, ascending) + at := i + if !ascending { + at = len(ids) - 1 - i + } + if err != nil || len(page.Turns) != 1 || page.Turns[0].ID != ids[at] { + t.Fatalf("page %d: %+v %v", i, page, err) + } + if (page.NextCursor != "") != (i < len(ids)-1) { + t.Fatalf("incorrect has_more: %+v", page) + } + cursor = page.Turns[0].ID + } + page, err := store.ListTurns(ctx, tenant, session, cursor, 1, ascending) + if err != nil || len(page.Turns) != 0 || page.Turns == nil || page.NextCursor != "" { + t.Fatalf("end: %+v %v", page, err) + } + } + otherTenant, otherSession, _ := newTurn(t, pool, sessions.TurnCancelled) + empty := uuid.New() + exec(t, pool, `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash) VALUES ($1, $2, 'codex', 'empty', 'hash')`, empty, tenantID) + for _, scope := range [][2]string{{text(otherTenant), session}, {tenant, text(otherSession)}, {tenant, uuid.NewString()}, {tenant, empty.String()}} { + if _, err := store.ListTurns(ctx, scope[0], scope[1], ids[0], 1, true); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("foreign cursor/session accepted: %v", err) + } + } + page, err := store.ListTurns(ctx, tenant, empty.String(), "", 20, false) + if err != nil || page.Turns == nil || len(page.Turns) != 0 { + t.Fatalf("empty session: %+v %v", page, err) + } + for _, limit := range []int{0, 101} { + if _, err := store.ListTurns(ctx, tenant, session, "", limit, false); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatalf("limit accepted: %v", err) + } + } +} diff --git a/services/core/internal/persistence/postgres/sessionpg/functions.go b/services/core/internal/persistence/postgres/sessionpg/functions.go index 8da77be69..8f82b4fdb 100644 --- a/services/core/internal/persistence/postgres/sessionpg/functions.go +++ b/services/core/internal/persistence/postgres/sessionpg/functions.go @@ -150,18 +150,3 @@ func (t *SessionTx) ApplyFunctionResult(ctx context.Context, turn, call string) } return t.q.ApplyFunctionResult(ctx, sqlc.ApplyFunctionResultParams{SessionID: t.session, TurnID: id, CallID: call}) } - -func (t *SessionTx) ApplyFunctionTurnStatus(ctx context.Context, turn, expected, status string) (sessions.Turn, error) { - id, err := parseID(turn) - if err != nil { - return sessions.Turn{}, err - } - row, err := t.q.TransitionTurn(ctx, sqlc.TransitionTurnParams{ID: id, SessionID: t.session, ExpectedStatus: expected, NewStatus: status, Outcome: []byte(`{}`)}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.Turn{}, sessions.ErrTurnConflict - } - if err != nil { - return sessions.Turn{}, err - } - return TurnFromRow(row), nil -} diff --git a/services/core/internal/persistence/postgres/sessionpg/model_execution.go b/services/core/internal/persistence/postgres/sessionpg/model_execution.go new file mode 100644 index 000000000..8f5cfbfbb --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/model_execution.go @@ -0,0 +1,54 @@ +package sessionpg + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + + "github.com/jackc/pgx/v5" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var _ sessions.ModelExecutionReader = (*Store)(nil) + +func (s *Store) SessionModelExecution(ctx context.Context, tenant, session string) (*v1.ModelProviderInput, error) { + tenantID, err := parseID(tenant) + if err != nil { + return nil, err + } + sessionID, err := parseID(session) + if err != nil { + return nil, err + } + ciphertext, err := s.units.Queries().GetSessionModelExecution(ctx, sqlc.GetSessionModelExecutionParams{TenantID: tenantID, SessionID: sessionID}) + if errors.Is(err, pgx.ErrNoRows) { + return nil, sessions.ErrNotFound + } + if err != nil { + return nil, err + } + if s.cipher == nil { + return nil, credentialcrypto.ErrUnavailable + } + raw, err := s.cipher.OpenModelExecution(ciphertext, tenant, session) + if err != nil { + return nil, fmt.Errorf("open session model execution: %w", err) + } + var provider v1.ModelProviderInput + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.DisallowUnknownFields() + if decoder.Decode(&provider) != nil || decoder.Decode(new(any)) != io.EOF { + return nil, errors.New("invalid stored model execution configuration") + } + if err := provider.Validate(); err != nil { + return nil, err + } + return &provider, nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/turn_reads.go b/services/core/internal/persistence/postgres/sessionpg/turn_reads.go new file mode 100644 index 000000000..8704de091 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/turn_reads.go @@ -0,0 +1,119 @@ +package sessionpg + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var _ sessions.TurnReader = (*Store)(nil) + +func (s *Store) GetTurn(ctx context.Context, tenant, session, turn string) (sessions.Turn, error) { + lookup, err := PublicTurnLookup(tenant, session, turn) + if err != nil { + return sessions.Turn{}, err + } + row, err := s.units.Queries().GetTurn(ctx, lookup) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.Turn{}, sessions.ErrNotFound + } + if err != nil { + return sessions.Turn{}, fmt.Errorf("get turn: %w", err) + } + return TurnFromRow(row), nil +} + +func (s *Store) ListTurns(ctx context.Context, tenant, session, cursor string, limit int, ascending bool) (sessions.TurnPage, error) { + if limit < 1 || limit > 100 { + return sessions.TurnPage{}, fmt.Errorf("%w: page size must be 1..100", sessions.ErrInvalidInput) + } + tenantID, err := parseID(tenant) + if err != nil { + return sessions.TurnPage{}, err + } + sessionID := pgunit.PathID(session) + q := s.units.Queries() + if _, err := q.GetSession(ctx, sqlc.GetSessionParams{TenantID: tenantID, ID: sessionID}); errors.Is(err, pgx.ErrNoRows) { + return sessions.TurnPage{}, sessions.ErrNotFound + } else if err != nil { + return sessions.TurnPage{}, fmt.Errorf("get session: %w", err) + } + params := sqlc.ListRootTurnsParams{TenantID: tenantID, SessionID: sessionID, PageLimit: int32(limit + 1), AfterID: pgtype.UUID{Valid: true}, Ascending: ascending} + if cursor != "" { + // A child Turn is not a Session Turn, so its ID is a missing cursor here. + after, err := s.GetTurn(ctx, tenant, session, pgunit.LookupCursor(cursor)) + if err != nil { + return sessions.TurnPage{}, err + } + params.AfterCreated = pgtype.Timestamptz{Time: after.CreatedAt, Valid: true} + params.AfterID, _ = parseID(after.ID) + } + rows, err := q.ListRootTurns(ctx, params) + if err != nil { + return sessions.TurnPage{}, fmt.Errorf("list turns: %w", err) + } + page := sessions.TurnPage{Turns: make([]sessions.Turn, 0, min(limit, len(rows)))} + if len(rows) > limit { + page.NextCursor = uuid.UUID(rows[limit-1].ID.Bytes).String() + rows = rows[:limit] + } + for _, row := range rows { + page.Turns = append(page.Turns, TurnFromRow(row)) + } + return page, nil +} + +func (s *Store) ListExecutionWork(ctx context.Context, after string, statuses, connectedDevices []string) ([]sessions.ExecutionWork, error) { + id, devices, err := ExecutionWorkCursor(after, connectedDevices) + if err != nil { + return nil, err + } + rows, err := s.units.Queries().ListExecutionWork(ctx, sqlc.ListExecutionWorkParams{AfterID: id, Statuses: statuses, ConnectedOnly: connectedDevices != nil, ConnectedDevices: devices}) + if err != nil { + return nil, err + } + work := make([]sessions.ExecutionWork, 0, len(rows)) + for _, row := range rows { + work = append(work, sessions.ExecutionWork{TenantID: uuid.UUID(row.TenantID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), TurnID: uuid.UUID(row.ID.Bytes).String(), Status: row.Status}) + } + return work, nil +} + +// ExecutionWorkCursor parses the cursor and connected devices of an execution +// work scan; an empty cursor starts at the first ID. A malformed ID is +// sessions.ErrInvalidInput. +func ExecutionWorkCursor(after string, connectedDevices []string) (pgtype.UUID, []pgtype.UUID, error) { + id := pgtype.UUID{Valid: true} + var err error + if after != "" { + id, err = parseID(after) + if err != nil { + return id, nil, err + } + } + devices := make([]pgtype.UUID, 0, len(connectedDevices)) + for _, value := range connectedDevices { + device, err := parseID(value) + if err != nil { + return id, nil, err + } + devices = append(devices, device) + } + return id, devices, nil +} + +// PublicTurnLookup resolves the caller's path identifiers of a Turn or a +// Turn-scoped resource. A malformed tenant is sessions.ErrInvalidInput; a +// malformed Session or Turn ID resolves as a missing one. +func PublicTurnLookup(tenant, session, turn string) (sqlc.GetTurnParams, error) { + id, err := parseID(tenant) + return sqlc.GetTurnParams{TenantID: id, SessionID: pgunit.PathID(session), ID: pgunit.PathID(turn)}, err +} diff --git a/services/core/internal/sandbox/providers/configuration_flow_test.go b/services/core/internal/sandbox/providers/configuration_flow_test.go index 6b1ac416f..b82a51051 100644 --- a/services/core/internal/sandbox/providers/configuration_flow_test.go +++ b/services/core/internal/sandbox/providers/configuration_flow_test.go @@ -136,7 +136,7 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { Sessions: s, SessionCreation: s, SessionEvents: s, - Turns: s, + Turns: struct{ api.Turns }{}, Items: struct{ api.Items }{}, Subagents: struct{ api.Subagents }{}, Artifacts: struct{ api.Artifacts }{}, diff --git a/services/core/internal/sessions/devices.go b/services/core/internal/sessions/devices.go index 417263183..d7fd40cf9 100644 --- a/services/core/internal/sessions/devices.go +++ b/services/core/internal/sessions/devices.go @@ -72,6 +72,14 @@ type DeviceReader interface { // ListEnrolledRuntimeBindings lists the enrolled user-managed Runtimes of // live Environments. ListEnrolledRuntimeBindings(ctx context.Context) ([]EnrolledRuntimeBinding, error) + // GetSessionExecutionBinding reads the Runtime device that executes the + // Session's Turns, with the native session that continues its history, + // once its Environment preparation completed; before that, and without an + // authorized bound device, it is ErrNotFound. + GetSessionExecutionBinding(ctx context.Context, tenant, session string) (ExecutionBinding, error) + // ListExecutionDevices lists the tenant's unrevoked devices that belong + // to no Environment, in ID order. + ListExecutionDevices(ctx context.Context, tenant string) ([]ExecutionDevice, error) } // DeviceStorage stores Runtime devices. diff --git a/services/core/internal/sessions/doc.go b/services/core/internal/sessions/doc.go index adb5fa21d..e86c58aab 100644 --- a/services/core/internal/sessions/doc.go +++ b/services/core/internal/sessions/doc.go @@ -7,12 +7,12 @@ // The Session writes that several operations share are procedures here, over // the transaction interfaces declared beside them: cancelling work, failing // and terminating an Environment, tracking input activity, the admission -// gates, admitting a function result, reading a Turn's required actions, -// appending to a Turn's execution journal and projecting its observations and -// admitted inputs into Items, Turn usage and Subagents. Service runs the pooled +// gates, moving a Turn's status, admitting a function result, reading a Turn's +// required actions, appending to a Turn's execution journal and projecting its +// observations and admitted inputs into Items, Turn usage and Subagents. Service runs the pooled // Session use cases, such as staging a Turn's Artifacts, over Storage, and // Reader declares the Session reads. ExecutionOperations runs the Session -// writes only the execution owner makes, such as recording function calls and -// their application receipts and journaling a Turn's observations, over the -// lease-bound ExecutionStorage. +// writes only the execution owner makes, such as moving and completing Turns, +// recording function calls and their application receipts and journaling a +// Turn's observations, over the lease-bound ExecutionStorage. package sessions diff --git a/services/core/internal/sessions/execution.go b/services/core/internal/sessions/execution.go index e222dc62c..db5eaaf13 100644 --- a/services/core/internal/sessions/execution.go +++ b/services/core/internal/sessions/execution.go @@ -20,5 +20,5 @@ func NewExecutionOperations(storage ExecutionStorage) (*ExecutionOperations, err type ExecutionStorage interface { EnvironmentExecution FunctionExecution - JournalExecution + TurnExecution } diff --git a/services/core/internal/sessions/execution_functions.go b/services/core/internal/sessions/execution_functions.go index aa6478c14..1285a1a9e 100644 --- a/services/core/internal/sessions/execution_functions.go +++ b/services/core/internal/sessions/execution_functions.go @@ -1,6 +1,9 @@ package sessions -import "context" +import ( + "context" + "encoding/json" +) // FunctionCallMatch is what a Turn's stored call shows about a reported call. type FunctionCallMatch struct { @@ -28,9 +31,9 @@ type FunctionTx interface { // ApplyFunctionResult records the application receipt of the call's // result. ApplyFunctionResult(ctx context.Context, turn, call string) error - // ApplyFunctionTurnStatus moves the Turn from expected to status and - // returns it as stored. A Turn no longer in expected is ErrTurnConflict. - ApplyFunctionTurnStatus(ctx context.Context, turn, expected, status string) (Turn, error) + // ApplyTurnStatus moves the Turn as change decides and returns it as + // stored. A Turn no longer in the expected status is ErrTurnConflict. + ApplyTurnStatus(ctx context.Context, turn string, change TurnStatusChange) (Turn, error) } // FunctionExecution is the lease-bound storage of function calls. @@ -150,7 +153,7 @@ func settleFunctionState(ctx context.Context, tx FunctionTx, turn Turn) error { status = TurnWaiting } if turn.Status != status { - if turn, err = tx.ApplyFunctionTurnStatus(ctx, turn.ID, turn.Status, status); err != nil { + if turn, err = tx.ApplyTurnStatus(ctx, turn.ID, TurnStatusChange{Expected: turn.Status, Status: status, Outcome: json.RawMessage(`{}`)}); err != nil { return err } if err := tx.AppendChanges(ctx, TurnChanges(turn, false)...); err != nil { diff --git a/services/core/internal/sessions/execution_functions_test.go b/services/core/internal/sessions/execution_functions_test.go index 42531197d..cc8ca606a 100644 --- a/services/core/internal/sessions/execution_functions_test.go +++ b/services/core/internal/sessions/execution_functions_test.go @@ -17,7 +17,7 @@ import ( type fakeExecutionStorage struct { t *testing.T withFunctionTurn func(ctx context.Context, tenant, session, turn string, apply func(context.Context, FunctionTx, Turn) error) error - withTurnJournal func(ctx context.Context, tenant, session string, apply func(context.Context, TurnJournalTx) error) error + withTurns func(ctx context.Context, tenant, session string, apply func(context.Context, TurnTx) error) error tx *fakeTx locked LockedSession @@ -40,12 +40,12 @@ func (f *fakeExecutionStorage) WithFunctionTurn(ctx context.Context, tenant, ses return f.withFunctionTurn(ctx, tenant, session, turn, apply) } -func (f *fakeExecutionStorage) WithTurnJournal(ctx context.Context, tenant, session string, apply func(context.Context, TurnJournalTx) error) error { +func (f *fakeExecutionStorage) WithTurns(ctx context.Context, tenant, session string, apply func(context.Context, TurnTx) error) error { f.t.Helper() - if f.withTurnJournal == nil { - f.t.Fatal("unexpected call to WithTurnJournal") + if f.withTurns == nil { + f.t.Fatal("unexpected call to WithTurns") } - return f.withTurnJournal(ctx, tenant, session, apply) + return f.withTurns(ctx, tenant, session, apply) } const testTenant = "4f0d0c35-8b8e-4d7c-9d1c-1f0a5b8a2e61" @@ -132,14 +132,15 @@ func TestRecordFunctionCall(t *testing.T) { t.Run("first call moves the Turn to waiting", func(t *testing.T) { running, waiting := turnWith(TurnInProgress), turnWith(TurnWaiting) var changes []SessionChange - f := &fakeFunctionTx{fakeTx: &fakeTx{t: t, loadUsage: returns(usage), appendChanges: collect(&changes)}, + moved := func(TurnStatusChange) (Turn, error) { return waiting, nil } + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t, loadUsage: returns(usage), appendChanges: collect(&changes), applyTurnStatus: moved}, matchFunctionCall: returns(FunctionCallMatch{}), createFunctionCall: done, - loadPendingFunctionCalls: returns([]FunctionCall{testCall}), applyFunctionTurnStatus: returns(waiting)} + loadPendingFunctionCalls: returns([]FunctionCall{testCall})} if err := functionOperations(t, f, running).RecordFunctionCall(t.Context(), testTenant, testSession, testTurn, testCall); err != nil { t.Fatal(err) } assertCalls(t, f.fakeTx, "MatchFunctionCall "+testTurn+" call", "CreateFunctionCall "+testTurn+" call", "LoadPendingFunctionCalls "+testTurn, - "ApplyFunctionTurnStatus "+testTurn+" in_progress waiting", "AppendChanges ", "LoadUsage", "AppendChanges agent.session.requires_action") + "ApplyTurnStatus "+testTurn+" in_progress waiting {}", "AppendChanges ", "LoadUsage", "AppendChanges agent.session.requires_action") actions, _ := LoadRequiredActions(t.Context(), &fakeFunctionTx{fakeTx: &fakeTx{t: t}, loadPendingFunctionCalls: returns([]FunctionCall{testCall})}, waiting) if !reflect.DeepEqual(changes, []SessionChange{ActivityChange(waiting, usage, actions)}) { t.Fatalf("changes %+v", changes) @@ -198,14 +199,15 @@ func TestConfirmFunctionResult(t *testing.T) { t.Run("last receipt resumes the Turn", func(t *testing.T) { waiting, running := turnWith(TurnWaiting), turnWith(TurnInProgress) var changes []SessionChange - f := &fakeFunctionTx{fakeTx: &fakeTx{t: t, loadUsage: returns(usage), appendChanges: collect(&changes)}, + moved := func(TurnStatusChange) (Turn, error) { return running, nil } + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t, loadUsage: returns(usage), appendChanges: collect(&changes), applyTurnStatus: moved}, loadFunctionCall: func() (FunctionCall, bool, error) { return submitted, true, nil }, applyFunctionResult: done, - loadPendingFunctionCalls: returns([]FunctionCall{}), applyFunctionTurnStatus: returns(running)} + loadPendingFunctionCalls: returns([]FunctionCall{})} if err := functionOperations(t, f, waiting).ConfirmFunctionResult(t.Context(), testTenant, testSession, testTurn, "call"); err != nil { t.Fatal(err) } assertCalls(t, f.fakeTx, "LoadFunctionCall "+testTurn+" call", "ApplyFunctionResult "+testTurn+" call", "LoadPendingFunctionCalls "+testTurn, - "ApplyFunctionTurnStatus "+testTurn+" waiting in_progress", "AppendChanges agent.session.turn.in_progress", "LoadUsage", "AppendChanges agent.session.in_progress") + "ApplyTurnStatus "+testTurn+" waiting in_progress {}", "AppendChanges agent.session.turn.in_progress", "LoadUsage", "AppendChanges agent.session.in_progress") want := append(TurnChanges(running, false), ActivityChange(running, usage, []v1.FunctionCallAction{})) if !reflect.DeepEqual(changes, want) { t.Fatalf("changes %+v", changes) diff --git a/services/core/internal/sessions/execution_journal.go b/services/core/internal/sessions/execution_journal.go deleted file mode 100644 index 95480f05d..000000000 --- a/services/core/internal/sessions/execution_journal.go +++ /dev/null @@ -1,30 +0,0 @@ -package sessions - -import "context" - -// JournalExecution records execution observations in Turn journals. -type JournalExecution interface { - // WithTurnJournal runs apply in a Session transaction on the lease-bound - // connection, with the tenant's Session locked, and passes apply the - // transaction's context. A deleted Session still records its Turns' - // observations. - WithTurnJournal(ctx context.Context, tenant, session string, apply func(context.Context, TurnJournalTx) error) error -} - -// AppendTurnEvents records an ordered batch of a Turn's execution observations -// in its journal from position first and projects them, as the -// AppendTurnEvents procedure decides. A malformed tenant, Session or Turn ID is -// ErrInvalidInput, before the batch is validated and before the Session is -// looked up. The journal is Core-internal; it is not the public event stream. -func (o *ExecutionOperations) AppendTurnEvents(ctx context.Context, tenant, session, turn string, first int32, events []ExecutionEvent) error { - if !validID(tenant) || !validID(session) { - return ErrInvalidInput - } - batch, err := NewJournalBatch(turn, first, events) - if err != nil { - return err - } - return o.storage.WithTurnJournal(ctx, tenant, session, func(ctx context.Context, tx TurnJournalTx) error { - return AppendTurnEvents(ctx, tx, batch) - }) -} diff --git a/services/core/internal/sessions/execution_turns.go b/services/core/internal/sessions/execution_turns.go new file mode 100644 index 000000000..a6fbf912c --- /dev/null +++ b/services/core/internal/sessions/execution_turns.go @@ -0,0 +1,306 @@ +package sessions + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" +) + +// maxOutcomeBytes bounds a Turn's terminal outcome. +const maxOutcomeBytes = 512 * 1024 + +// TurnStatusChange is a decided change of a Turn's status. +type TurnStatusChange struct { + Expected string + Status string + Outcome json.RawMessage + // SourceCompletedAt is when the native source completed a Turn that ends + // completed, zero when it reported none. + SourceCompletedAt time.Time +} + +// TurnTransitionTx is the Session transaction TransitionTurn runs in. +type TurnTransitionTx interface { + ProjectionTx + ComputeAdmissionTx + // LoadTurn reads one of the Session's Turns; a missing one is ErrNotFound. + LoadTurn(ctx context.Context, turn string) (Turn, error) + // ApplyTurnStatus moves the Turn as change decides and returns it as + // stored. The database clock stamps the first start of a Turn that moves + // to in progress and the end of a Turn that ends, unless the change carries + // the source completion time. A Turn no longer in the expected status is + // ErrTurnConflict. + ApplyTurnStatus(ctx context.Context, turn string, change TurnStatusChange) (Turn, error) + // LoadEnding reads the facts EndTurn settles an ended Turn with. + LoadEnding(ctx context.Context, turn string) (Ending, error) + // ApplyTurnEnd writes what EndTurn decided for an ended Turn. + ApplyTurnEnd(ctx context.Context, turn string, end TurnEnd) error +} + +// TurnTx is the Session transaction of the Turn execution operations. +type TurnTx interface { + TurnTransitionTx + TurnEventTx + TurnJournalTx + // HasUnappliedInputs reports whether the Turn has message inputs after + // sequence appliedThrough. + HasUnappliedInputs(ctx context.Context, turn string, appliedThrough int64) (bool, error) + // RememberNativeSession records the native session that continues the + // Session's history on its bound device. A Session without a bound device + // is ErrNotFound. + RememberNativeSession(ctx context.Context, native string) error + // BeginArtifactCapture marks that the Turn captures its Artifacts. A Turn + // that already does is ErrTurnConflict. + BeginArtifactCapture(ctx context.Context, turn string) error +} + +// TurnExecution is the lease-bound storage of the Turn execution operations. +type TurnExecution interface { + // WithTurns runs apply in one transaction on the execution lease, under + // the tenant's Session lock, and commits only when apply succeeds. A + // malformed ID is ErrInvalidInput and a missing Session ErrNotFound; a + // publicly deleted Session still journals and settles its Turns. + WithTurns(ctx context.Context, tenant, session string, apply func(context.Context, TurnTx) error) error +} + +// decideTransition decides the status change of current that transition asks +// for. A transition ValidTransition does not allow, an outcome that is not one +// JSON object of at most 512 KiB, and an outcome on a Turn that does not end +// are ErrInvalidInput. A Turn no longer in the expected status, and a Turn +// with a cancellation request that would resume, are ErrTurnConflict. +func decideTransition(current Turn, transition TurnTransition) (TurnStatusChange, error) { + if !ValidTransition(transition.ExpectedStatus, transition.Status) || len(transition.Outcome) > maxOutcomeBytes { + return TurnStatusChange{}, fmt.Errorf("%w: invalid turn transition or outcome size", ErrInvalidInput) + } + outcome, err := jsonobject.Normalize(transition.Outcome) + if err != nil { + return TurnStatusChange{}, fmt.Errorf("%w: %w", ErrInvalidInput, err) + } + if !TerminalStatus(transition.Status) && string(outcome) != "{}" { + return TurnStatusChange{}, fmt.Errorf("%w: outcome requires a terminal status", ErrInvalidInput) + } + if current.Status != transition.ExpectedStatus || (transition.Status == TurnInProgress && !current.CancelRequestedAt.IsZero()) { + return TurnStatusChange{}, ErrTurnConflict + } + return TurnStatusChange{Expected: transition.ExpectedStatus, Status: transition.Status, Outcome: outcome}, nil +} + +// TransitionTurn moves the Session's root Turn as transition asks, a +// compare-and-set on its status, and publishes what the move settles. A Turn +// that keeps running journals its event. A Turn that ends projects its +// outcome, then settles as EndTurn decides from the Turn read after the +// projection, which may record its usage. An ended Turn is never reopened and +// its outcome never overwritten, including by retries. +func TransitionTurn(ctx context.Context, tx TurnTransitionTx, turn string, transition TurnTransition) (Turn, error) { + current, err := tx.LoadTurn(ctx, turn) + if err != nil { + return Turn{}, err + } + change, err := decideTransition(current, transition) + if err != nil { + return Turn{}, err + } + if change.Expected == TurnQueued && change.Status == TurnInProgress { + if err := CheckComputeAdmission(ctx, tx); err != nil { + return Turn{}, err + } + } + moved, err := tx.ApplyTurnStatus(ctx, turn, change) + if err != nil { + return Turn{}, err + } + if !TerminalStatus(moved.Status) { + if err := tx.AppendChanges(ctx, TurnChanges(moved, false)...); err != nil { + return Turn{}, err + } + return moved, nil + } + if err := ProjectSource(ctx, tx, Source{Turn: moved.ID, Kind: "execution_" + moved.Status, Payload: moved.Outcome, CreatedAt: moved.CompletedAt}); err != nil { + return Turn{}, err + } + return settleEnd(ctx, tx, turn) +} + +// settleEnd settles a Turn that ended as EndTurn decides, from the Turn as the +// transaction has left it, and returns that Turn. +func settleEnd(ctx context.Context, tx TurnTransitionTx, turn string) (Turn, error) { + ended, err := tx.LoadTurn(ctx, turn) + if err != nil { + return Turn{}, err + } + ending, err := tx.LoadEnding(ctx, turn) + if err != nil { + return Turn{}, err + } + return ended, tx.ApplyTurnEnd(ctx, turn, EndTurn(ended, ending)) +} + +// decideCompletion decides whether an execution may end a Turn in status +// current with status: an in-progress Turn ends in any terminal status, and a +// waiting Turn, which still waits on actions, only fails or is cancelled. +// Anything else is ErrTurnConflict. +func decideCompletion(current, status string) error { + if current == TurnInProgress || (current == TurnWaiting && status != TurnCompleted) { + return nil + } + return ErrTurnConflict +} + +// sourceCompletedAt reads the native completion time a completed outcome +// reports in done.source_completed_at_ms, zero when it reports none. An +// outcome that does not decode, or a time that is not positive, is +// ErrInvalidInput. Native and Core timestamps come from independent host +// clocks: the Turn keeps the source time, while committed activity uses the +// database clock. +func sourceCompletedAt(outcome json.RawMessage) (time.Time, error) { + var snapshot struct { + Done *struct { + SourceCompletedAtMS *int64 `json:"source_completed_at_ms"` + } `json:"done"` + } + if json.Unmarshal(outcome, &snapshot) != nil { + return time.Time{}, ErrInvalidInput + } + if snapshot.Done == nil || snapshot.Done.SourceCompletedAtMS == nil { + return time.Time{}, nil + } + ms := *snapshot.Done.SourceCompletedAtMS + if ms <= 0 { + return time.Time{}, ErrInvalidInput + } + return time.UnixMilli(ms), nil +} + +// checkInputsApplied requires that the execution applied every message input +// of the Turn through sequence appliedThrough; otherwise it is +// ErrUnappliedInputs. +func checkInputsApplied(ctx context.Context, tx TurnTx, turn string, appliedThrough int64) error { + pending, err := tx.HasUnappliedInputs(ctx, turn, appliedThrough) + if err != nil { + return err + } + if pending { + return ErrUnappliedInputs + } + return nil +} + +// TransitionTurn moves the tenant's root Turn as transition asks, as the +// TransitionTurn procedure decides. A dispatcher claims a queued Turn as in +// progress before it sends the Turn's work to a Runtime. A malformed ID is +// ErrInvalidInput. +func (o *ExecutionOperations) TransitionTurn(ctx context.Context, tenant, session, turn string, transition TurnTransition) (Turn, error) { + if !validID(tenant) || !validID(session) || !validID(turn) { + return Turn{}, ErrInvalidInput + } + var result Turn + err := o.storage.WithTurns(ctx, tenant, session, func(ctx context.Context, tx TurnTx) error { + var err error + result, err = TransitionTurn(ctx, tx, turn, transition) + return err + }) + return result, err +} + +// CompleteExecution ends the Turn an execution ran with the execution's +// terminal status and outcome. The status change, the outcome's journal entry, +// the native session that continues the Session's history and the Turn's +// settlement, which publishes or discards its staged Artifacts, commit +// together. A Turn completes only when the execution applied its message +// inputs through appliedThrough; otherwise it is ErrUnappliedInputs. A +// malformed ID, a status that does not end the Turn, an outcome that is not +// one JSON object of at most 512 KiB, a native session ID over 512 bytes and a +// negative appliedThrough are ErrInvalidInput. +func (o *ExecutionOperations) CompleteExecution(ctx context.Context, tenant, session, turn, status string, outcome json.RawMessage, nativeID string, appliedThrough int64) (Turn, error) { + if !validID(tenant) || !validID(session) || !validID(turn) { + return Turn{}, ErrInvalidInput + } + if !TerminalStatus(status) || len(outcome) > maxOutcomeBytes || len(nativeID) > 512 || appliedThrough < 0 { + return Turn{}, ErrInvalidInput + } + outcome, err := jsonobject.Normalize(outcome) + if err != nil { + return Turn{}, fmt.Errorf("%w: %w", ErrInvalidInput, err) + } + var result Turn + err = o.storage.WithTurns(ctx, tenant, session, func(ctx context.Context, tx TurnTx) error { + current, found, err := tx.LoadJournalTurn(ctx, turn) + if err != nil { + return err + } + if !found { + return ErrNotFound + } + if err := decideCompletion(current.Status, status); err != nil { + return err + } + change := TurnStatusChange{Expected: current.Status, Status: status, Outcome: outcome} + if status == TurnCompleted { + if err := checkInputsApplied(ctx, tx, turn, appliedThrough); err != nil { + return err + } + if change.SourceCompletedAt, err = sourceCompletedAt(outcome); err != nil { + return err + } + } + if _, err := tx.ApplyTurnStatus(ctx, turn, change); err != nil { + return err + } + if err := AppendTurnEvent(ctx, tx, turn, current.EventCount, ExecutionEvent{Kind: "execution_" + status, Payload: outcome}); err != nil { + return err + } + if nativeID != "" { + if err := tx.RememberNativeSession(ctx, nativeID); err != nil { + return err + } + } + result, err = settleEnd(ctx, tx, turn) + return err + }) + return result, err +} + +// BeginTurnArtifactCapture admits the capture of an in-progress Turn's +// Artifacts once the execution applied its message inputs through +// appliedThrough; otherwise it is ErrUnappliedInputs. A Turn that is not in +// progress, has a cancellation request or already captures is ErrTurnConflict. +// A malformed ID or a negative appliedThrough is ErrInvalidInput. +func (o *ExecutionOperations) BeginTurnArtifactCapture(ctx context.Context, tenant, session, turn string, appliedThrough int64) error { + if !validID(tenant) || !validID(session) || !validID(turn) || appliedThrough < 0 { + return ErrInvalidInput + } + return o.storage.WithTurns(ctx, tenant, session, func(ctx context.Context, tx TurnTx) error { + current, err := tx.LoadTurn(ctx, turn) + if err != nil { + return err + } + if current.Status != TurnInProgress || !current.CancelRequestedAt.IsZero() { + return ErrTurnConflict + } + if err := checkInputsApplied(ctx, tx, turn, appliedThrough); err != nil { + return err + } + return tx.BeginArtifactCapture(ctx, turn) + }) +} + +// AppendTurnEvents records an ordered batch of a Turn's execution observations +// in its journal from position first and projects them, as the +// AppendTurnEvents procedure decides. A malformed tenant, Session or Turn ID is +// ErrInvalidInput, before the batch is validated and before the Session is +// looked up. The journal is Core-internal; it is not the public event stream. +func (o *ExecutionOperations) AppendTurnEvents(ctx context.Context, tenant, session, turn string, first int32, events []ExecutionEvent) error { + if !validID(tenant) || !validID(session) { + return ErrInvalidInput + } + batch, err := NewJournalBatch(turn, first, events) + if err != nil { + return err + } + return o.storage.WithTurns(ctx, tenant, session, func(ctx context.Context, tx TurnTx) error { + return AppendTurnEvents(ctx, tx, batch) + }) +} diff --git a/services/core/internal/sessions/execution_turns_test.go b/services/core/internal/sessions/execution_turns_test.go new file mode 100644 index 000000000..88fbca53f --- /dev/null +++ b/services/core/internal/sessions/execution_turns_test.go @@ -0,0 +1,294 @@ +package sessions + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "reflect" + "strings" + "testing" + "time" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" +) + +func TestDecideTransition(t *testing.T) { + cancelling := turnWith(TurnWaiting) + cancelling.CancelRequestedAt = time.Unix(1, 0) + big := json.RawMessage(`{"x":"` + strings.Repeat("x", maxOutcomeBytes) + `"}`) + for name, test := range map[string]struct { + current Turn + transition TurnTransition + want TurnStatusChange + err error + }{ + "start is admitted": { + current: turnWith(TurnQueued), transition: TurnTransition{ExpectedStatus: TurnQueued, Status: TurnInProgress}, + want: TurnStatusChange{Expected: TurnQueued, Status: TurnInProgress, Outcome: json.RawMessage(`{}`)}, + }, + "wait keeps running": { + current: turnWith(TurnInProgress), transition: TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnWaiting, Outcome: json.RawMessage(`{}`)}, + want: TurnStatusChange{Expected: TurnInProgress, Status: TurnWaiting, Outcome: json.RawMessage(`{}`)}, + }, + "failure ends": { + current: turnWith(TurnInProgress), transition: TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnFailed, Outcome: json.RawMessage(` {"error":"x"} `)}, + want: TurnStatusChange{Expected: TurnInProgress, Status: TurnFailed, Outcome: json.RawMessage(`{"error":"x"}`)}, + }, + "cancelled waiting Turn ends": {current: cancelling, transition: TurnTransition{ExpectedStatus: TurnWaiting, Status: TurnCancelled}, want: TurnStatusChange{Expected: TurnWaiting, Status: TurnCancelled, Outcome: json.RawMessage(`{}`)}}, + "ended Turn never moves": {current: turnWith(TurnCompleted), transition: TurnTransition{ExpectedStatus: TurnCompleted, Status: TurnFailed}, err: ErrInvalidInput}, + "skipped status": {current: turnWith(TurnQueued), transition: TurnTransition{ExpectedStatus: TurnQueued, Status: TurnCompleted}, err: ErrInvalidInput}, + "oversized outcome": {current: turnWith(TurnInProgress), transition: TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnFailed, Outcome: big}, err: ErrInvalidInput}, + "outcome not an object": {current: turnWith(TurnInProgress), transition: TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnFailed, Outcome: json.RawMessage(`[]`)}, err: ErrInvalidInput}, + "outcome on a running Turn": {current: turnWith(TurnInProgress), transition: TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnWaiting, Outcome: json.RawMessage(`{"a":1}`)}, err: ErrInvalidInput}, + "invalid before conflict": {current: turnWith(TurnCompleted), transition: TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnWaiting, Outcome: json.RawMessage(`{"a":1}`)}, err: ErrInvalidInput}, + "status moved": {current: turnWith(TurnCompleted), transition: TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnFailed}, err: ErrTurnConflict}, + "cancellation blocks resuming": {current: cancelling, transition: TurnTransition{ExpectedStatus: TurnWaiting, Status: TurnInProgress}, err: ErrTurnConflict}, + "cancellation blocks the start": {current: func() Turn { turn := turnWith(TurnQueued); turn.CancelRequestedAt = time.Unix(1, 0); return turn }(), transition: TurnTransition{ExpectedStatus: TurnQueued, Status: TurnInProgress}, err: ErrTurnConflict}, + } { + got, err := decideTransition(test.current, test.transition) + if !errors.Is(err, test.err) { + t.Fatalf("%s: %v, want %v", name, err, test.err) + } + if !reflect.DeepEqual(got, test.want) { + t.Fatalf("%s: %+v, want %+v", name, got, test.want) + } + } +} + +func TestDecideCompletion(t *testing.T) { + for _, test := range []struct { + current, status string + err error + }{ + {TurnInProgress, TurnCompleted, nil}, + {TurnInProgress, TurnFailed, nil}, + {TurnInProgress, TurnCancelled, nil}, + {TurnWaiting, TurnFailed, nil}, + {TurnWaiting, TurnCancelled, nil}, + {TurnWaiting, TurnCompleted, ErrTurnConflict}, + {TurnQueued, TurnFailed, ErrTurnConflict}, + {TurnCompleted, TurnFailed, ErrTurnConflict}, + } { + if err := decideCompletion(test.current, test.status); !errors.Is(err, test.err) { + t.Fatalf("%s to %s: %v, want %v", test.current, test.status, err, test.err) + } + } +} + +func TestSourceCompletedAt(t *testing.T) { + for outcome, want := range map[string]struct { + at time.Time + err error + }{ + `{}`: {}, + `{"done":{}}`: {}, + `{"done":{"source_completed_at_ms":1500}}`: {at: time.UnixMilli(1500)}, + `{"done":{"source_completed_at_ms":0}}`: {err: ErrInvalidInput}, + `{"done":{"source_completed_at_ms":"x"}}`: {err: ErrInvalidInput}, + } { + at, err := sourceCompletedAt(json.RawMessage(outcome)) + if !errors.Is(err, want.err) || !at.Equal(want.at) { + t.Fatalf("%s: %v %v, want %v %v", outcome, at, err, want.at, want.err) + } + } +} + +// moves is a fake ApplyTurnStatus that returns the Turn in its new status. +func moves(change TurnStatusChange) (Turn, error) { + turn := turnWith(change.Status) + turn.Outcome = change.Outcome + return turn, nil +} + +// A queued Turn starts only after compute admission, and a Turn that keeps +// running journals its event. +func TestTransitionTurnStarts(t *testing.T) { + f := &fakeTx{t: t, loadTurn: returns(turnWith(TurnQueued)), loadComputeSuspension: returns(false), applyTurnStatus: moves, appendChanges: func([]SessionChange) error { return nil }} + turn, err := TransitionTurn(t.Context(), f, testTurn, TurnTransition{ExpectedStatus: TurnQueued, Status: TurnInProgress}) + if err != nil || turn.Status != TurnInProgress { + t.Fatal(turn, err) + } + assertCalls(t, f, "LoadTurn "+testTurn, "LoadComputeSuspension", "ApplyTurnStatus "+testTurn+" queued in_progress {}", "AppendChanges agent.session.turn.in_progress") + + f = &fakeTx{t: t, loadTurn: returns(turnWith(TurnQueued)), loadComputeSuspension: returns(true)} + if _, err := TransitionTurn(t.Context(), f, testTurn, TurnTransition{ExpectedStatus: TurnQueued, Status: TurnInProgress}); !errors.Is(err, ErrTurnConflict) { + t.Fatal(err) + } + assertCalls(t, f, "LoadTurn "+testTurn, "LoadComputeSuspension") +} + +// A Turn that ends projects its outcome, then settles from the Turn as the +// projection left it. +func TestTransitionTurnEnds(t *testing.T) { + outcome := `{"done":{"usage":{"tokens":{"cached_input_tokens":0,"input_tokens":1,"output_tokens":2,"reasoning_output_tokens":0,"total_tokens":3}}}}` + settled := turnWith(TurnFailed) + settled.Usage = json.RawMessage(`{"total_tokens":3}`) + read := []Turn{turnWith(TurnInProgress), settled} + var end TurnEnd + f := &fakeTx{ + t: t, loadTurn: func() (Turn, error) { turn := read[0]; read = read[1:]; return turn, nil }, applyTurnStatus: moves, + putTurnUsage: func(v1.TokenUsage) error { return nil }, loadEnding: returns(Ending{}), applyTurnEnd: func(e TurnEnd) error { end = e; return nil }, + } + turn, err := TransitionTurn(t.Context(), f, testTurn, TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnFailed, Outcome: json.RawMessage(outcome)}) + if err != nil { + t.Fatal(err) + } + assertCalls(t, f, "LoadTurn "+testTurn, "ApplyTurnStatus "+testTurn+" in_progress failed "+outcome, "PutTurnUsage "+testTurn, "LoadTurn "+testTurn, "LoadEnding "+testTurn, "ApplyTurnEnd "+testTurn) + if string(turn.Usage) != string(settled.Usage) || end.Artifacts != DiscardArtifacts { + t.Fatal(turn, end) + } +} + +// A rejected transition writes nothing. +func TestTransitionTurnRejectsWithoutWrites(t *testing.T) { + f := &fakeTx{t: t, loadTurn: returns(turnWith(TurnCompleted))} + if _, err := TransitionTurn(t.Context(), f, testTurn, TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnFailed}); !errors.Is(err, ErrTurnConflict) { + t.Fatal(err) + } + assertCalls(t, f, "LoadTurn "+testTurn) + + f = &fakeTx{t: t, loadTurn: func() (Turn, error) { return Turn{}, ErrNotFound }} + if _, err := TransitionTurn(t.Context(), f, testTurn, TurnTransition{ExpectedStatus: TurnInProgress, Status: TurnFailed}); !errors.Is(err, ErrNotFound) { + t.Fatal(err) + } + assertCalls(t, f, "LoadTurn "+testTurn) +} + +// turnOperations runs every Turn transaction of the tenant's Session on tx. +func turnOperations(t *testing.T, tx *fakeTx) *ExecutionOperations { + t.Helper() + operations, err := NewExecutionOperations(&fakeExecutionStorage{t: t, withTurns: func(ctx context.Context, tenant, session string, apply func(context.Context, TurnTx) error) error { + if tenant != testTenant || session != testSession { + t.Fatalf("Session %s %s", tenant, session) + } + return apply(ctx, tx) + }}) + if err != nil { + t.Fatal(err) + } + return operations +} + +// A completed execution checks its inputs, moves the Turn with the source +// completion time, journals the outcome after the last entry, remembers the +// native session and settles the Turn, which publishes its Artifacts. +func TestCompleteExecutionOrder(t *testing.T) { + outcome := `{"done":{"source_completed_at_ms":1500}}` + var end TurnEnd + f := &fakeTx{ + t: t, loadJournalTurn: func() (JournalTurn, bool, error) { + return JournalTurn{Status: TurnInProgress, EventCount: 4}, true, nil + }, + hasUnappliedInputs: returns(false), applyTurnStatus: moves, insertEvent: done, countEvents: done, loadEventSources: returns([]Source(nil)), + rememberNativeSession: done, loadTurn: returns(turnWith(TurnCompleted)), loadEnding: returns(Ending{}), + applyTurnEnd: func(e TurnEnd) error { end = e; return nil }, + } + turn, err := turnOperations(t, f).CompleteExecution(t.Context(), testTenant, testSession, testTurn, TurnCompleted, json.RawMessage(outcome), "native-1", 7) + if err != nil || turn.Status != TurnCompleted { + t.Fatal(turn, err) + } + assertCalls(t, f, + "LoadJournalTurn "+testTurn, "HasUnappliedInputs "+testTurn+" 7", + "ApplyTurnStatus "+testTurn+" in_progress completed "+outcome+" 1500", + "InsertEvent "+testTurn+" 5 execution_completed", "CountEvents "+testTurn+" 1 "+fmt.Sprint(len(outcome)), "LoadEventSources "+testTurn+" 5", + "RememberNativeSession native-1", "LoadTurn "+testTurn, "LoadEnding "+testTurn, "ApplyTurnEnd "+testTurn, + ) + if end.Artifacts != PublishArtifacts { + t.Fatal(end) + } +} + +// A failed execution of a waiting Turn skips the input check and the native +// session, and settles the Turn, which discards its Artifacts. +func TestCompleteExecutionFails(t *testing.T) { + var end TurnEnd + f := &fakeTx{ + t: t, loadJournalTurn: func() (JournalTurn, bool, error) { return JournalTurn{Status: TurnWaiting}, true, nil }, + applyTurnStatus: moves, insertEvent: done, countEvents: done, loadEventSources: returns([]Source(nil)), + loadTurn: returns(turnWith(TurnFailed)), loadEnding: returns(Ending{}), applyTurnEnd: func(e TurnEnd) error { end = e; return nil }, + } + if _, err := turnOperations(t, f).CompleteExecution(t.Context(), testTenant, testSession, testTurn, TurnFailed, nil, "", 0); err != nil { + t.Fatal(err) + } + assertCalls(t, f, "LoadJournalTurn "+testTurn, "ApplyTurnStatus "+testTurn+" waiting failed {}", "InsertEvent "+testTurn+" 1 execution_failed", "CountEvents "+testTurn+" 1 2", "LoadEventSources "+testTurn+" 1", "LoadTurn "+testTurn, "LoadEnding "+testTurn, "ApplyTurnEnd "+testTurn) + if end.Artifacts != DiscardArtifacts { + t.Fatal(end) + } +} + +func TestCompleteExecutionRejects(t *testing.T) { + for name, test := range map[string]struct { + turn JournalTurn + found bool + pending bool + outcome string + err error + calls []string + }{ + "missing Turn": {err: ErrNotFound, calls: []string{"LoadJournalTurn " + testTurn}}, + "waiting Turn completes": {turn: JournalTurn{Status: TurnWaiting}, found: true, err: ErrTurnConflict, calls: []string{"LoadJournalTurn " + testTurn}}, + "unapplied inputs": {turn: JournalTurn{Status: TurnInProgress}, found: true, pending: true, err: ErrUnappliedInputs, calls: []string{"LoadJournalTurn " + testTurn, "HasUnappliedInputs " + testTurn + " 0"}}, + "bad source time": {turn: JournalTurn{Status: TurnInProgress}, found: true, outcome: `{"done":{"source_completed_at_ms":-1}}`, err: ErrInvalidInput, calls: []string{"LoadJournalTurn " + testTurn, "HasUnappliedInputs " + testTurn + " 0"}}, + } { + f := &fakeTx{t: t, loadJournalTurn: func() (JournalTurn, bool, error) { return test.turn, test.found, nil }, hasUnappliedInputs: returns(test.pending)} + if _, err := turnOperations(t, f).CompleteExecution(t.Context(), testTenant, testSession, testTurn, TurnCompleted, json.RawMessage(test.outcome), "", 0); !errors.Is(err, test.err) { + t.Fatalf("%s: %v", name, err) + } + assertCalls(t, f, test.calls...) + } + for name, call := range map[string]func(*ExecutionOperations) error{ + "malformed Turn": func(o *ExecutionOperations) error { + _, err := o.CompleteExecution(t.Context(), testTenant, testSession, "x", TurnFailed, nil, "", 0) + return err + }, + "running status": func(o *ExecutionOperations) error { + _, err := o.CompleteExecution(t.Context(), testTenant, testSession, testTurn, TurnWaiting, nil, "", 0) + return err + }, + "long native": func(o *ExecutionOperations) error { + _, err := o.CompleteExecution(t.Context(), testTenant, testSession, testTurn, TurnFailed, nil, strings.Repeat("n", 513), 0) + return err + }, + "negative applied": func(o *ExecutionOperations) error { + _, err := o.CompleteExecution(t.Context(), testTenant, testSession, testTurn, TurnFailed, nil, "", -1) + return err + }, + "array outcome": func(o *ExecutionOperations) error { + _, err := o.CompleteExecution(t.Context(), testTenant, testSession, testTurn, TurnFailed, json.RawMessage(`[]`), "", 0) + return err + }, + } { + if err := call(unusedStorage(t)); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("%s: %v", name, err) + } + } +} + +func TestBeginTurnArtifactCapture(t *testing.T) { + f := &fakeTx{t: t, loadTurn: returns(turnWith(TurnInProgress)), hasUnappliedInputs: returns(false), beginArtifactCapture: done} + if err := turnOperations(t, f).BeginTurnArtifactCapture(t.Context(), testTenant, testSession, testTurn, 3); err != nil { + t.Fatal(err) + } + assertCalls(t, f, "LoadTurn "+testTurn, "HasUnappliedInputs "+testTurn+" 3", "BeginArtifactCapture "+testTurn) + + cancelling := turnWith(TurnInProgress) + cancelling.CancelRequestedAt = time.Unix(1, 0) + for _, current := range []Turn{turnWith(TurnWaiting), cancelling} { + f = &fakeTx{t: t, loadTurn: returns(current)} + if err := turnOperations(t, f).BeginTurnArtifactCapture(t.Context(), testTenant, testSession, testTurn, 3); !errors.Is(err, ErrTurnConflict) { + t.Fatal(current.Status, err) + } + assertCalls(t, f, "LoadTurn "+testTurn) + } + + f = &fakeTx{t: t, loadTurn: returns(turnWith(TurnInProgress)), hasUnappliedInputs: returns(true)} + if err := turnOperations(t, f).BeginTurnArtifactCapture(t.Context(), testTenant, testSession, testTurn, 3); !errors.Is(err, ErrUnappliedInputs) { + t.Fatal(err) + } + assertCalls(t, f, "LoadTurn "+testTurn, "HasUnappliedInputs "+testTurn+" 3") + + if err := unusedStorage(t).BeginTurnArtifactCapture(t.Context(), testTenant, testSession, testTurn, -1); !errors.Is(err, ErrInvalidInput) { + t.Fatal(err) + } +} diff --git a/services/core/internal/sessions/functions_test.go b/services/core/internal/sessions/functions_test.go index 1b1a33f5b..76d661868 100644 --- a/services/core/internal/sessions/functions_test.go +++ b/services/core/internal/sessions/functions_test.go @@ -24,7 +24,6 @@ type fakeFunctionTx struct { createFunctionCall func() error loadFunctionCall func() (FunctionCall, bool, error) applyFunctionResult func() error - applyFunctionTurnStatus func() (Turn, error) } var ( @@ -77,11 +76,6 @@ func (f *fakeFunctionTx) ApplyFunctionResult(_ context.Context, turn, call strin return f.applyFunctionResult() } -func (f *fakeFunctionTx) ApplyFunctionTurnStatus(_ context.Context, turn, expected, status string) (Turn, error) { - f.record("ApplyFunctionTurnStatus", f.applyFunctionTurnStatus != nil, turn, expected, status) - return f.applyFunctionTurnStatus() -} - func cancelling(status string) Turn { turn := turnWith(status) turn.CancelRequestedAt = time.Unix(1700000001, 0) diff --git a/services/core/internal/sessions/journal_test.go b/services/core/internal/sessions/journal_test.go index 40c0e08fe..35152b1da 100644 --- a/services/core/internal/sessions/journal_test.go +++ b/services/core/internal/sessions/journal_test.go @@ -139,7 +139,7 @@ func TestAppendTurnEvent(t *testing.T) { func TestExecutionOperationsValidationOrder(t *testing.T) { big := json.RawMessage(`{"text":"` + strings.Repeat("x", 400*1024) + `"}`) oversized := []ExecutionEvent{{Kind: "delta", Payload: big}, {Kind: "delta", Payload: big}, {Kind: "delta", Payload: big}} - operations, err := NewExecutionOperations(&fakeExecutionStorage{t: t, withTurnJournal: func(context.Context, string, string, func(context.Context, TurnJournalTx) error) error { + operations, err := NewExecutionOperations(&fakeExecutionStorage{t: t, withTurns: func(context.Context, string, string, func(context.Context, TurnTx) error) error { return ErrNotFound }}) if err != nil { diff --git a/services/core/internal/sessions/model_execution.go b/services/core/internal/sessions/model_execution.go new file mode 100644 index 000000000..632ac47bb --- /dev/null +++ b/services/core/internal/sessions/model_execution.go @@ -0,0 +1,16 @@ +package sessions + +import ( + "context" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" +) + +// ModelExecutionReader reads the model execution a Session froze at creation. +type ModelExecutionReader interface { + // SessionModelExecution opens the model provider the tenant's Session + // froze at creation. A Session without one is ErrNotFound, and a service + // without the credential key credentialcrypto.ErrUnavailable; a frozen + // provider that does not open, decode or validate is an internal error. + SessionModelExecution(ctx context.Context, tenant, session string) (*v1.ModelProviderInput, error) +} diff --git a/services/core/internal/sessions/reader.go b/services/core/internal/sessions/reader.go index fe80400aa..32dd253bd 100644 --- a/services/core/internal/sessions/reader.go +++ b/services/core/internal/sessions/reader.go @@ -10,5 +10,7 @@ type Reader interface { EnvironmentReader ExecutorCredentialReader ItemReader + ModelExecutionReader SubagentReader + TurnReader } diff --git a/services/core/internal/sessions/transaction_test.go b/services/core/internal/sessions/transaction_test.go index 2f54fb6f2..6cc0809bd 100644 --- a/services/core/internal/sessions/transaction_test.go +++ b/services/core/internal/sessions/transaction_test.go @@ -89,6 +89,11 @@ type fakeTx struct { recordTerminalActivity func() error loadChildItem func() (StoredChildItem, bool, error) putChildItem func(ChildItem) error + + applyTurnStatus func(TurnStatusChange) (Turn, error) + hasUnappliedInputs func() (bool, error) + rememberNativeSession func() error + beginArtifactCapture func() error } var ( @@ -96,8 +101,6 @@ var ( _ InputStartTx = (*fakeTx)(nil) _ ComputeAdmissionTx = (*fakeTx)(nil) _ EnvironmentDeviceTx = (*fakeTx)(nil) - _ TurnJournalTx = (*fakeTx)(nil) - _ TurnEventTx = (*fakeTx)(nil) _ InputProjectionTx = (*fakeTx)(nil) _ InitializationTx = (*fakeTx)(nil) _ ConnectionTx = (*fakeTx)(nil) @@ -105,6 +108,7 @@ var ( _ EnrollmentTx = (*fakeTx)(nil) _ FileWriteReservationTx = (*fakeTx)(nil) _ FileWriteSettlementTx = (*fakeTx)(nil) + _ TurnTx = (*fakeTx)(nil) _ EnvironmentExecutorCredentialTx = (*fakeTx)(nil) ) @@ -312,6 +316,30 @@ func (f *fakeTx) PutChildItem(_ context.Context, item ChildItem) error { return f.putChildItem(item) } +func (f *fakeTx) ApplyTurnStatus(_ context.Context, turn string, change TurnStatusChange) (Turn, error) { + detail := []string{turn, change.Expected, change.Status, string(change.Outcome)} + if !change.SourceCompletedAt.IsZero() { + detail = append(detail, fmt.Sprint(change.SourceCompletedAt.UnixMilli())) + } + f.record("ApplyTurnStatus", f.applyTurnStatus != nil, detail...) + return f.applyTurnStatus(change) +} + +func (f *fakeTx) HasUnappliedInputs(_ context.Context, turn string, appliedThrough int64) (bool, error) { + f.record("HasUnappliedInputs", f.hasUnappliedInputs != nil, turn, fmt.Sprint(appliedThrough)) + return f.hasUnappliedInputs() +} + +func (f *fakeTx) RememberNativeSession(_ context.Context, native string) error { + f.record("RememberNativeSession", f.rememberNativeSession != nil, native) + return f.rememberNativeSession() +} + +func (f *fakeTx) BeginArtifactCapture(_ context.Context, turn string) error { + f.record("BeginArtifactCapture", f.beginArtifactCapture != nil, turn) + return f.beginArtifactCapture() +} + // returns is a fake method that reads value. func returns[T any](value T) func() (T, error) { return func() (T, error) { return value, nil } diff --git a/services/core/internal/sessions/turn.go b/services/core/internal/sessions/turn.go index c7ad7e11c..f6a4574c5 100644 --- a/services/core/internal/sessions/turn.go +++ b/services/core/internal/sessions/turn.go @@ -1,6 +1,7 @@ package sessions import ( + "context" "encoding/json" "time" ) @@ -53,6 +54,8 @@ func ValidTransition(from, to string) bool { } } +// TurnTransition asks to move a Turn from ExpectedStatus to Status. Only a +// Turn that ends carries an outcome. type TurnTransition struct { ExpectedStatus string Status string @@ -91,3 +94,22 @@ type ExecutionBinding struct { NativeSessionID string HasStartedTurn bool } + +// TurnReader reads root Turns. Subagent Turns are not Session Turns; +// SubagentReader reads them. +type TurnReader interface { + // GetTurn reads the tenant's root Turn of the Session. A malformed tenant + // is ErrInvalidInput; a malformed Session or Turn ID, a Subagent Turn ID + // and a missing Turn are ErrNotFound. + GetTurn(ctx context.Context, tenant, session, turn string) (Turn, error) + // ListTurns pages the root Turns of the tenant's Session that was not + // publicly deleted by creation time and ID, after the Turn cursor names. + // A page size outside 1..100 is ErrInvalidInput; a missing Session or + // cursor Turn is ErrNotFound. + ListTurns(ctx context.Context, tenant, session, cursor string, limit int, ascending bool) (TurnPage, error) + // ListExecutionWork lists, in ID order after the Turn after names, up to + // 100 root Turns in statuses across tenants; the queued Turns of a publicly + // deleted Session are not work. With connectedDevices not nil, it lists + // only the Turns of tenants that own one of those devices, unrevoked. + ListExecutionWork(ctx context.Context, after string, statuses, connectedDevices []string) ([]ExecutionWork, error) +} diff --git a/services/core/internal/store/admin_session_archive_test.go b/services/core/internal/store/admin_session_archive_test.go index ba8d933d1..511a6f458 100644 --- a/services/core/internal/store/admin_session_archive_test.go +++ b/services/core/internal/store/admin_session_archive_test.go @@ -222,7 +222,7 @@ func TestManagedSessionArchiveAuditFailureRollsBack(t *testing.T) { if _, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1); err != nil { t.Fatal(err) } - turn, err := s.GetTurn(t.Context(), tenant, session.ID, input.TurnID) + turn, err := sessionAdapter(s).GetTurn(t.Context(), tenant, session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnInProgress || turn.CancelRequestedAt.IsZero() { t.Fatal("archive did not request cancellation or fabricated settlement", turn, err) } diff --git a/services/core/internal/store/admin_session_archive_worker_http_test.go b/services/core/internal/store/admin_session_archive_worker_http_test.go index 146de71bd..df4b565e1 100644 --- a/services/core/internal/store/admin_session_archive_worker_http_test.go +++ b/services/core/internal/store/admin_session_archive_worker_http_test.go @@ -119,7 +119,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { if err != nil || allocation.ID != owner.ID || allocation.State != "cleanup_pending" { t.Fatal("archive did not retain cleanup ownership", allocation, err) } - turn, err := s.GetTurn(t.Context(), project.TenantID, active.ID, input.TurnID) + turn, err := store.SessionAdapter(s).GetTurn(t.Context(), project.TenantID, active.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnCancelled { t.Fatal("archive did not cancel queued work", turn, err) } diff --git a/services/core/internal/store/agent_execution_defaults_http_test.go b/services/core/internal/store/agent_execution_defaults_http_test.go index 67eb2f2da..95b9ead57 100644 --- a/services/core/internal/store/agent_execution_defaults_http_test.go +++ b/services/core/internal/store/agent_execution_defaults_http_test.go @@ -82,7 +82,7 @@ func TestAgentExecutionDefaultsPublicSnapshotAndPrecedence(t *testing.T) { if json.Unmarshal(session.Configuration, &cfg) != nil || cfg.Agent.Model != model { t.Fatal("model snapshot changed") } - provider, err := st.SessionModelExecution(t.Context(), tenant, sessionID) + provider, err := store.SessionAdapter(st).SessionModelExecution(t.Context(), tenant, sessionID) if err != nil || provider.BaseURL != endpoint || provider.APIKey != key { t.Fatal("provider snapshot mismatch", err) } @@ -106,7 +106,7 @@ func TestAgentExecutionDefaultsPublicSnapshotAndPrecedence(t *testing.T) { created := id(call("POST", "/v1/agents/sessions", nativeHarness, uuid.NewString(), 201)) assertSnapshot(created, "model-override", "https://override.example/v1", "override-canary") resolved, err := st.GetSession(t.Context(), tenant, created) - frozen, providerErr := st.SessionModelExecution(t.Context(), tenant, created) + frozen, providerErr := store.SessionAdapter(st).SessionModelExecution(t.Context(), tenant, created) if err != nil || providerErr != nil || resolved.Engine != "claude_sdk" || frozen == nil || frozen.Protocol != "anthropic" { t.Fatal("harness override did not freeze the native provider", err, providerErr) } diff --git a/services/core/internal/store/archive_cancellation_test.go b/services/core/internal/store/archive_cancellation_test.go index 3434dd0cb..4c173ad2b 100644 --- a/services/core/internal/store/archive_cancellation_test.go +++ b/services/core/internal/store/archive_cancellation_test.go @@ -182,7 +182,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { t.Fatal(err) } - current, err := s.GetTurn(t.Context(), h.tenant, session.ID, input.TurnID) + current, err := store.SessionAdapter(s).GetTurn(t.Context(), h.tenant, session.ID, input.TurnID) if err != nil || current.Status != sessions.TurnWaiting || current.CancelRequestedAt.IsZero() { t.Fatal("archive must request rather than invent cancellation", current, err) } diff --git a/services/core/internal/store/artifact_lifecycle.go b/services/core/internal/store/artifact_lifecycle.go deleted file mode 100644 index eb244e2fa..000000000 --- a/services/core/internal/store/artifact_lifecycle.go +++ /dev/null @@ -1,48 +0,0 @@ -package store - -import ( - "context" - "errors" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -// BeginTurnArtifactCapture separates native completion from bounded output publication. -// Later messages use the existing reservation path instead of the finished executor. -func (s *Store) BeginTurnArtifactCapture(ctx context.Context, tenantID, sessionID, turnID string, appliedThrough int64) error { - lookup, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) - if err != nil { - return err - } - if appliedThrough < 0 { - return sessions.ErrInvalidInput - } - return s.withSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - turn, err := q.GetTurn(ctx, lookup) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - if turn.Status != sessions.TurnInProgress || turn.CancelRequestedAt.Valid { - return sessions.ErrTurnConflict - } - pending, err := q.HasUnappliedMessages(ctx, sqlc.HasUnappliedMessagesParams{SessionID: session, TurnID: lookup.ID, Sequence: appliedThrough}) - if err != nil { - return err - } - if pending { - return sessions.ErrUnappliedInputs - } - count, err := q.BeginTurnArtifactCapture(ctx, sqlc.BeginTurnArtifactCaptureParams{SessionID: session, ID: lookup.ID}) - if err == nil && count != 1 { - return sessions.ErrTurnConflict - } - return err - }) -} diff --git a/services/core/internal/store/claude_execution_test.go b/services/core/internal/store/claude_execution_test.go index 3a94c3829..b76101ae7 100644 --- a/services/core/internal/store/claude_execution_test.go +++ b/services/core/internal/store/claude_execution_test.go @@ -68,7 +68,7 @@ func TestClaudeWorkerSelectsStoredEngineAndRestrictiveCapabilities(t *testing.T) }() queued := func() { time.Sleep(650 * time.Millisecond) - turn, err := h.s.GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnQueued { t.Fatal(turn, err) } @@ -91,7 +91,7 @@ func TestClaudeWorkerSelectsStoredEngineAndRestrictiveCapabilities(t *testing.T) } h.write(input.TurnID, proto.TypeDone, proto.DonePayload{Content: "done", Metadata: map[string]any{proto.DoneMetaAgentSessionID: "claude-native"}}) waitTurn(t, h, input.TurnID, sessions.TurnCompleted) - bound, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || bound.NativeSessionID != "claude-native" { t.Fatal(bound, err) } @@ -112,7 +112,7 @@ func TestClaudeDispatcherRejectsUnsupportedConfigurationBeforeClaim(t *testing.T if result := <-h.run(t.Context(), input.TurnID); result.err == nil { t.Fatal("unsupported configuration claimed") } - turn, err := h.s.GetTurn(t.Context(), h.tenant, h.session.ID, input.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnQueued { t.Fatal(turn, err) } @@ -126,7 +126,7 @@ func TestClaudeInvalidImageResultRejectsWholeBatchBeforePersistence(t *testing.T worker := startOwnedWorker(t, t.Context(), h.db, h.d, h.owner()) defer func() { ctx, cancel := context.WithCancel(context.Background()); cancel(); _ = worker.Run(ctx) }() input := h.message("start", "Run") - if _, err := h.s.TransitionTurn(t.Context(), h.tenant, h.session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err := store.TransitionTurn(t.Context(), h.s, h.tenant, h.session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } call := sessions.FunctionCall{CallID: "public-call", ExecutorCallID: "native-call", Name: "lookup_ticket", Arguments: json.RawMessage(`{"ticket":"42"}`)} @@ -148,7 +148,7 @@ func TestClaudeInvalidImageResultRejectsWholeBatchBeforePersistence(t *testing.T if err != nil || saved.Result != nil || saved.Applied { t.Fatal(saved, err) } - turn, err := h.s.GetTurn(t.Context(), h.tenant, h.session.ID, input.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnWaiting || !turn.CancelRequestedAt.IsZero() { t.Fatal(turn, err) } diff --git a/services/core/internal/store/claude_mcp_test.go b/services/core/internal/store/claude_mcp_test.go index b72536d98..0a81a5a6c 100644 --- a/services/core/internal/store/claude_mcp_test.go +++ b/services/core/internal/store/claude_mcp_test.go @@ -11,6 +11,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestClaudeMCPWaitsForCapableRuntime(t *testing.T) { @@ -55,7 +56,7 @@ func TestClaudeMCPWaitsForCapableRuntime(t *testing.T) { } }() time.Sleep(650 * time.Millisecond) - turn, err := h.s.GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnQueued { t.Fatal("incapable runtime claimed work", turn, err) } @@ -121,7 +122,7 @@ func TestClaudeMCPUnsupportedSnapshotRejectedBeforeClaim(t *testing.T) { if result := <-h.run(t.Context(), input.TurnID); result.err == nil { t.Fatal("unsupported MCP configuration claimed") } - turn, err := h.s.GetTurn(t.Context(), h.tenant, h.session.ID, input.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnQueued { t.Fatal(turn, err) } diff --git a/services/core/internal/store/command_output_test.go b/services/core/internal/store/command_output_test.go index 8ae0b4ba5..b8e90e586 100644 --- a/services/core/internal/store/command_output_test.go +++ b/services/core/internal/store/command_output_test.go @@ -26,7 +26,7 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err = s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err = store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } event := func(kind, raw string) sessions.ExecutionEvent { @@ -70,7 +70,7 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { if err := journal.AppendTurnEvents(ctx, tenant, session.ID, input.TurnID, 4, final); err != nil { t.Fatal(err) } - if _, err := s.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{}`), "", input.Sequence); err != nil { + if _, err := journal.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{}`), "", input.Sequence); err != nil { t.Fatal(err) } // Reopening the Store recovers committed Items without creating events. diff --git a/services/core/internal/store/deployment_model_providers_http_test.go b/services/core/internal/store/deployment_model_providers_http_test.go index e3a09f6c6..63866bca3 100644 --- a/services/core/internal/store/deployment_model_providers_http_test.go +++ b/services/core/internal/store/deployment_model_providers_http_test.go @@ -83,7 +83,7 @@ func TestDeploymentModelProvidersHTTP(t *testing.T) { } providerOf := func(sessionID string) string { t.Helper() - provider, err := st.SessionModelExecution(t.Context(), tenant, sessionID) + provider, err := store.SessionAdapter(st).SessionModelExecution(t.Context(), tenant, sessionID) if err != nil { t.Fatal(err) } @@ -335,7 +335,7 @@ func TestNoneSessionRetryAfterDeploymentDefaultChanges(t *testing.T) { } key := uuid.NewString() original := create(key) - if provider, err := st.SessionModelExecution(t.Context(), tenant, original); err != nil || provider.APIKey != "first-default-key" { + if provider, err := store.SessionAdapter(st).SessionModelExecution(t.Context(), tenant, original); err != nil || provider.APIKey != "first-default-key" { t.Fatal("none Session did not freeze the deployment default", err) } setDefault("rotated-default-key") @@ -348,7 +348,7 @@ func TestNoneSessionRetryAfterDeploymentDefaultChanges(t *testing.T) { if create(key) != original { t.Fatal("retry after removal created another Session") } - if provider, err := st.SessionModelExecution(t.Context(), tenant, original); err != nil || provider.APIKey != "first-default-key" { + if provider, err := store.SessionAdapter(st).SessionModelExecution(t.Context(), tenant, original); err != nil || provider.APIKey != "first-default-key" { t.Fatal("retry changed the frozen provider", err) } var revision uuid.UUID @@ -401,7 +401,7 @@ func TestDeploymentProviderResolutionPairsRevisionDuringReplacement(t *testing.T if err = db.pool.QueryRow(t.Context(), "SELECT deployment_provider_revision FROM session_execution_configuration WHERE session_id=$1", session.ID).Scan(&revision); err != nil || revision != original.Revision { t.Fatal("tuple revision changed", err) } - frozen, err := st.SessionModelExecution(t.Context(), tenant, session.ID) + frozen, err := store.SessionAdapter(st).SessionModelExecution(t.Context(), tenant, session.ID) if err != nil || frozen == nil || *frozen != provider { t.Fatal("tuple bundle changed", err) } diff --git a/services/core/internal/store/devices.go b/services/core/internal/store/devices.go deleted file mode 100644 index f61266d6e..000000000 --- a/services/core/internal/store/devices.go +++ /dev/null @@ -1,44 +0,0 @@ -package store - -import ( - "context" - "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/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -func (s *Store) GetSessionExecutionBinding(ctx context.Context, tenantID, sessionID string) (sessions.ExecutionBinding, error) { - params, err := sessionpg.DeviceLookup(tenantID, sessionID) - if err != nil { - return sessions.ExecutionBinding{}, err - } - if err := s.requireInitializedEnvironment(ctx, params.TenantID, params.ID); err != nil { - return sessions.ExecutionBinding{}, err - } - row, err := s.queries.GetSessionExecutionBinding(ctx, sqlc.GetSessionExecutionBindingParams(params)) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ExecutionBinding{}, sessions.ErrNotFound - } - if err != nil { - return sessions.ExecutionBinding{}, err - } - return sessions.ExecutionBinding{ - Device: executionDevice(row.ID, row.Name, row.EnvironmentID), - NativeSessionID: row.NativeSessionID, - HasStartedTurn: row.HasStartedTurn, - }, nil -} - -func executionDevice(id pgtype.UUID, name string, environmentID pgtype.UUID) sessions.ExecutionDevice { - value := sessions.ExecutionDevice{ID: uuid.UUID(id.Bytes).String(), Name: name} - if environmentID.Valid { - value.EnvironmentID = uuid.UUID(environmentID.Bytes).String() - } - return value -} diff --git a/services/core/internal/store/dispatch_test.go b/services/core/internal/store/dispatch_test.go index 112b42b3a..e5ca9514c 100644 --- a/services/core/internal/store/dispatch_test.go +++ b/services/core/internal/store/dispatch_test.go @@ -268,7 +268,7 @@ func TestExecutionDispatchSteeringAndNativeContinuity(t *testing.T) { newStore, db := newTestStoreDB(t) h.s, h.db, h.owned = newStore, db, nil h.d.Store = newStore - bound, err := newStore.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) + bound, err := store.SessionAdapter(newStore).GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || bound.NativeSessionID != "native-thread-1" { t.Fatalf("native binding lost: %+v %v", bound, err) } @@ -300,7 +300,7 @@ func TestExecutionCancellationRequiresReceiptAndSurvivesContextEnd(t *testing.T) if cancel.DeliveryID == "" { t.Fatal("cancellation has no receipt identity") } - current, err := h.s.GetTurn(context.Background(), h.tenant, h.session.ID, first.TurnID) + current, err := store.SessionAdapter(h.s).GetTurn(context.Background(), h.tenant, h.session.ID, first.TurnID) if err != nil || current.Status != sessions.TurnInProgress { t.Fatal("cancel finished before receipt") } @@ -379,15 +379,16 @@ func TestExecutionOutcomeAndNativeBindingCommitTogether(t *testing.T) { h := newDispatchHarness(t) first := h.message("first", "Run") ctx := context.Background() - _, err := h.s.TransitionTurn(ctx, h.tenant, h.session.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err := store.TransitionTurn(ctx, h.s, h.tenant, h.session.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } late := h.message("second", "Late") - if _, err := h.s.CompleteExecution(ctx, h.tenant, h.session.ID, first.TurnID, sessions.TurnCompleted, []byte(`{}`), "native-one", first.Sequence); !errors.Is(err, sessions.ErrUnappliedInputs) { + operations := h.owner().Sessions + if _, err := operations.CompleteExecution(ctx, h.tenant, h.session.ID, first.TurnID, sessions.TurnCompleted, []byte(`{}`), "native-one", first.Sequence); !errors.Is(err, sessions.ErrUnappliedInputs) { t.Fatalf("unapplied completion: %v", err) } - bound, _ := h.s.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) + bound, _ := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if bound.NativeSessionID != "" { t.Fatal("native ID committed without outcome") } @@ -397,7 +398,7 @@ func TestExecutionOutcomeAndNativeBindingCommitTogether(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - _, err := h.s.CompleteExecution(ctx, h.tenant, h.session.ID, first.TurnID, sessions.TurnCompleted, []byte(`{}`), native, late.Sequence) + _, err := operations.CompleteExecution(ctx, h.tenant, h.session.ID, first.TurnID, sessions.TurnCompleted, []byte(`{}`), native, late.Sequence) errs <- err }() } @@ -448,7 +449,7 @@ func TestExecutionRejectsRuntimeMissingCapabilityBeforeClaim(t *testing.T) { if _, err := h.bound().Run(context.Background(), h.tenant, h.session.ID, first.TurnID); err == nil || err.Error() != tc.message { t.Fatalf("Run error = %v, want %q", err, tc.message) } - turn, err := h.s.GetTurn(context.Background(), h.tenant, h.session.ID, first.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(context.Background(), h.tenant, h.session.ID, first.TurnID) if err != nil || turn.Status != sessions.TurnQueued { t.Fatal("Runtime without a required capability claimed work") } diff --git a/services/core/internal/store/environment_claim_worker_test.go b/services/core/internal/store/environment_claim_worker_test.go index c4dd8249c..8efa54b68 100644 --- a/services/core/internal/store/environment_claim_worker_test.go +++ b/services/core/internal/store/environment_claim_worker_test.go @@ -10,6 +10,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { @@ -32,7 +33,7 @@ func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { t.Fatal(err) } } - turn, err := s.GetTurn(t.Context(), tenant, pending.SessionID, turnID) + turn, err := store.SessionAdapter(s).GetTurn(t.Context(), tenant, pending.SessionID, turnID) if err != nil || turn.Status != sessions.TurnInProgress || (deleted && turn.CancelRequestedAt.IsZero()) { t.Fatal("promotion did not retain the active claim", turn, err) } @@ -50,7 +51,7 @@ func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { t.Fatal(err) } awaitRelease() - turn, err = s.GetTurn(t.Context(), tenant, pending.SessionID, turnID) + turn, err = store.SessionAdapter(s).GetTurn(t.Context(), tenant, pending.SessionID, turnID) var outcome struct { ErrorCode string `json:"error_code"` } diff --git a/services/core/internal/store/environment_directory_test.go b/services/core/internal/store/environment_directory_test.go index 087c97f60..d98cf67a4 100644 --- a/services/core/internal/store/environment_directory_test.go +++ b/services/core/internal/store/environment_directory_test.go @@ -10,6 +10,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -133,7 +134,7 @@ func TestEnvironmentDirectoryWorkerReadsWithoutExecutionPrerequisites(t *testing if err != nil || session.LastTurn != nil || session.EnvironmentInputActivity != nil { t.Fatal("directory read manufactured execution") } - bound, err := h.s.GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) if err != nil || bound.Device.ID != h.device.ID || bound.NativeSessionID != "" { t.Fatal("directory read changed native history identity") } diff --git a/services/core/internal/store/environment_expiry_dispatch_test.go b/services/core/internal/store/environment_expiry_dispatch_test.go index 9482a8dea..28770e2f1 100644 --- a/services/core/internal/store/environment_expiry_dispatch_test.go +++ b/services/core/internal/store/environment_expiry_dispatch_test.go @@ -35,7 +35,7 @@ func TestWorkerEnvironmentExpiryAtFullExecutionCapacity(t *testing.T) { makeEnvironmentExpiryDue(t, pool, &due) waitEnvironmentExpiry(t, h.s, tenant, due) for i, request := range requests { - turn, err := h.s.GetTurn(t.Context(), h.tenant, active[i].ID, request.ID) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, active[i].ID, request.ID) if err != nil || turn.Status != sessions.TurnInProgress { t.Fatal("expiry was not observed at full capacity", turn, err) } diff --git a/services/core/internal/store/environment_input_activity_test.go b/services/core/internal/store/environment_input_activity_test.go index 1ed22d4ec..af7f32f8d 100644 --- a/services/core/internal/store/environment_input_activity_test.go +++ b/services/core/internal/store/environment_input_activity_test.go @@ -107,7 +107,7 @@ func TestEnvironmentInputActivitySettlementAndNewerWork(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := writer.TransitionTurn(t.Context(), tenant, session.ID, prior[0].TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed, Outcome: []byte(`{}`)}); err != nil { + if _, err := transitionTurn(t.Context(), writer, tenant, session.ID, prior[0].TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed, Outcome: []byte(`{}`)}); err != nil { t.Fatal(err) } reservation := reserveEnvironmentInput(t, s, tenant, session.ID, "waiting") diff --git a/services/core/internal/store/environment_input_claim_test.go b/services/core/internal/store/environment_input_claim_test.go index af872bf5b..4f588dad2 100644 --- a/services/core/internal/store/environment_input_claim_test.go +++ b/services/core/internal/store/environment_input_claim_test.go @@ -53,7 +53,7 @@ func TestEnvironmentInputConcurrentPromotionClaimsOnce(t *testing.T) { if received != count || fresh != 1 { t.Fatal("promotion authorized multiple starts", received, fresh) } - turn, err := s.GetTurn(t.Context(), tenant, session.ID, turnID) + turn, err := sessionAdapter(s).GetTurn(t.Context(), tenant, session.ID, turnID) if err != nil || turn.Status != sessions.TurnInProgress || turn.StartedAt.IsZero() { t.Fatal("promotion did not persist its execution claim", turn, err) } diff --git a/services/core/internal/store/environment_inputs.go b/services/core/internal/store/environment_inputs.go index 5fbfedec1..c2d2d4b7f 100644 --- a/services/core/internal/store/environment_inputs.go +++ b/services/core/internal/store/environment_inputs.go @@ -232,13 +232,12 @@ func settleEnvironmentInput(ctx context.Context, q *sqlc.Queries, tenantID strin return sessions.EnvironmentInputReservation{}, err } if state == sessions.EnvironmentInputAdmitted { - params, err := sessionpg.TurnLookup(tenantID, result.SessionID, result.Receipts[0].TurnID) + tenant, err := parseID(tenantID) if err != nil { return sessions.EnvironmentInputReservation{}, err } - if _, err := transitionTurn(ctx, q, params, sessions.TurnTransition{ - ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress, Outcome: json.RawMessage(`{}`), - }); err != nil { + start := sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress} + if _, err := sessions.TransitionTurn(ctx, sessionpg.BindSession(q, tenant, row.SessionID), result.Receipts[0].TurnID, start); err != nil { return sessions.EnvironmentInputReservation{}, err } } diff --git a/services/core/internal/store/environment_setup_test.go b/services/core/internal/store/environment_setup_test.go index 04c5e9456..b30dec563 100644 --- a/services/core/internal/store/environment_setup_test.go +++ b/services/core/internal/store/environment_setup_test.go @@ -46,7 +46,7 @@ func TestEnvironmentSetupEncryptedSnapshotAndIsolation(t *testing.T) { if _, err := s.CreateSession(t.Context(), tenant, input); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal("changed setup retried", err) } - if _, err := s.GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("uninitialized execution admitted", err) } } diff --git a/services/core/internal/store/environment_steering_order_test.go b/services/core/internal/store/environment_steering_order_test.go index cff36bdbe..94cfcf094 100644 --- a/services/core/internal/store/environment_steering_order_test.go +++ b/services/core/internal/store/environment_steering_order_test.go @@ -56,7 +56,7 @@ func TestEnvironmentActiveInputSerializesWithCompletion(t *testing.T) { admitted <- admission{value, err} } complete := func() { - _, err := writer.CompleteExecution(ctx, tenant, session.ID, original.TurnID, sessions.TurnCompleted, nil, "", original.Sequence) + _, err := completeExecution(ctx, t, writer, tenant, session.ID, original.TurnID, sessions.TurnCompleted, nil, "", original.Sequence) completed <- err } first, second := input, complete @@ -87,7 +87,7 @@ func TestEnvironmentActiveInputSerializesWithCompletion(t *testing.T) { if err != nil || retry.ID != got.value.ID || !retry.Deadline.Equal(got.value.Deadline) || len(retry.Receipts) != 2 || !retry.Receipts[0].Replayed || retry.Receipts[0].TurnID != prepared.Receipts[0].TurnID { t.Fatal("active retry replaced its original reservation", retry, err) } - if _, err := writer.CompleteExecution(ctx, tenant, session.ID, prepared.Receipts[0].TurnID, sessions.TurnCompleted, nil, "", prepared.Receipts[1].Sequence); err != nil { + if _, err := completeExecution(ctx, t, writer, tenant, session.ID, prepared.Receipts[0].TurnID, sessions.TurnCompleted, nil, "", prepared.Receipts[1].Sequence); err != nil { t.Fatal(err) } } else { @@ -95,7 +95,7 @@ func TestEnvironmentActiveInputSerializesWithCompletion(t *testing.T) { t.Fatal("admitted input escaped the original Turn or application fence", completionErr, got.value) } environmentInputHistory(t, pool, session.ID, 1, 3) - if _, err := writer.CompleteExecution(ctx, tenant, session.ID, original.TurnID, sessions.TurnCompleted, nil, "", got.value.Receipts[1].Sequence); err != nil { + if _, err := completeExecution(ctx, t, writer, tenant, session.ID, original.TurnID, sessions.TurnCompleted, nil, "", got.value.Receipts[1].Sequence); err != nil { t.Fatal("completion after controlled application failed", err) } } diff --git a/services/core/internal/store/environment_worker_helpers_test.go b/services/core/internal/store/environment_worker_helpers_test.go index e437e2f0e..8e6dcfefd 100644 --- a/services/core/internal/store/environment_worker_helpers_test.go +++ b/services/core/internal/store/environment_worker_helpers_test.go @@ -122,7 +122,7 @@ func awaitWorkerEnvironmentRun(t *testing.T, ctx context.Context, s *store.Store if len(run.Reservation.Receipts) == 0 { return false } - run.Turn, err = s.GetTurn(ctx, tenant, pending.SessionID, run.Reservation.Receipts[0].TurnID) + run.Turn, err = store.SessionAdapter(s).GetTurn(ctx, tenant, pending.SessionID, run.Reservation.Receipts[0].TurnID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/execution_events_test.go b/services/core/internal/store/execution_events_test.go index 43b39d69f..2ba74d5c9 100644 --- a/services/core/internal/store/execution_events_test.go +++ b/services/core/internal/store/execution_events_test.go @@ -94,7 +94,7 @@ func TestExecutionDoesNotCompleteAfterEventPersistenceFailure(t *testing.T) { if outcome.ErrorCode != "event_persistence_failed" { t.Fatal(outcome.ErrorCode) } - bound, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || bound.NativeSessionID != "failed-native" || outcome.Done.Usage.InputTokens != 13 { t.Fatalf("terminal failure lost native continuity or usage: %+v %+v %v", bound, outcome, err) } diff --git a/services/core/internal/store/execution_test.go b/services/core/internal/store/execution_test.go index 4b0c2f5ba..f2675a51b 100644 --- a/services/core/internal/store/execution_test.go +++ b/services/core/internal/store/execution_test.go @@ -11,8 +11,10 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" "github.com/jackc/pgx/v5/pgxpool" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" @@ -43,6 +45,42 @@ func sessionExecution(t *testing.T, lease *pgunit.Lease) *sessions.ExecutionOper return operations } +// transitionTurn moves the Turn as the execution owner does, in a Session +// transaction on s's writer: the lease of a NewExecution writer, otherwise the +// pool. +func transitionTurn(ctx context.Context, s *Store, tenant, session, turn string, transition sessions.TurnTransition) (sessions.Turn, error) { + var moved sessions.Turn + err := s.withLockedSession(ctx, tenant, session, false, func(ctx context.Context, q *sqlc.Queries, id pgtype.UUID, _ sessions.LockedSession) error { + owner, err := parseID(tenant) + if err != nil { + return err + } + moved, err = sessions.TransitionTurn(ctx, sessionpg.BindSession(q, owner, id), turn, transition) + return err + }) + return moved, err +} + +// completeExecution completes the Turn's execution through the Session +// execution operations on s's lease or, for a pooled s, on the execution lease +// it holds for the call. +func completeExecution(ctx context.Context, t testing.TB, s *Store, tenant, session, turn, status string, outcome json.RawMessage, native string, appliedThrough int64) (sessions.Turn, error) { + t.Helper() + lease := s.lease + if lease == nil { + var err error + if lease, err = pgunit.AcquireLease(ctx, s.pool); err != nil { + t.Fatal(err) + } + defer func() { _ = lease.Close(context.Background()) }() + } + operations, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) + if err != nil { + t.Fatal(err) + } + return operations.CompleteExecution(ctx, tenant, session, turn, status, outcome, native, appliedThrough) +} + // functionExecution builds the Session execution operations on the execution // lease of a pool of its own, so the test may close its Store's pool. func functionExecution(t *testing.T) *sessions.ExecutionOperations { @@ -122,14 +160,14 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { } } mustReject("binding", sessionExecution(t, writer.lease).BindSessionDevice(t.Context(), tenant, queued.ID, host.ID)) - _, err = writer.TransitionTurn(t.Context(), tenant, queued.ID, pending.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = operations.TransitionTurn(t.Context(), tenant, queued.ID, pending.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) mustReject("claim", err) mustReject("journal", operations.AppendTurnEvents(t.Context(), tenant, active.ID, input.TurnID, 1, []sessions.ExecutionEvent{{Kind: "delta", Payload: json.RawMessage(`{"delta":"stale"}`)}})) mustReject("callback", operations.RecordFunctionCall(t.Context(), tenant, active.ID, input.TurnID, functionCallFixture("late"))) mustReject("receipt", operations.ConfirmFunctionResult(t.Context(), tenant, waiting.ID, waitInput.TurnID, call.CallID)) - _, err = writer.CompleteExecution(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnCompleted, json.RawMessage(`{"done":{"content":"stale"}}`), "stale-native", input.Sequence) + _, err = operations.CompleteExecution(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnCompleted, json.RawMessage(`{"done":{"content":"stale"}}`), "stale-native", input.Sequence) mustReject("completion", err) - _, err = writer.TransitionTurn(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed}) + _, err = operations.TransitionTurn(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed}) mustReject("reconciliation", err) _, err = writer.ExpireEnvironmentInputs(t.Context()) mustReject("input expiry", err) @@ -153,7 +191,7 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { if _, err = sessionAdapter(s).GetSessionDevice(t.Context(), tenant, queued.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("stale binding persisted", err) } - queuedTurn, err := s.GetTurn(t.Context(), tenant, queued.ID, pending.TurnID) + queuedTurn, err := sessionAdapter(s).GetTurn(t.Context(), tenant, queued.ID, pending.TurnID) if err != nil || queuedTurn.Status != sessions.TurnQueued { t.Fatal("queued work changed", queuedTurn, err) } @@ -162,14 +200,15 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { if err = successor.lease.CheckOwnership(t.Context()); err != nil { t.Fatal(err) } - if _, err = successor.CompleteExecution(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnCompleted, json.RawMessage(`{"done":{"content":"accepted"}}`), "successor-native", input.Sequence); err != nil { + successorOperations := sessionExecution(t, successor.lease) + if _, err = successorOperations.CompleteExecution(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnCompleted, json.RawMessage(`{"done":{"content":"accepted"}}`), "successor-native", input.Sequence); err != nil { t.Fatal(err) } - bound, err := s.GetSessionExecutionBinding(t.Context(), tenant, active.ID) + bound, err := sessionAdapter(s).GetSessionExecutionBinding(t.Context(), tenant, active.ID) if err != nil || bound.NativeSessionID != "successor-native" { t.Fatal(bound, err) } - _, err = successor.TransitionTurn(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed}) + _, err = successorOperations.TransitionTurn(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed}) if !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("terminal CAS changed", err) } @@ -207,7 +246,7 @@ func TestExecutionWriterSerializesWritesOnItsLease(t *testing.T) { results := make(chan error, len(tasks)*2) for _, task := range tasks { group.Go(func() { - _, err := writer.TransitionTurn(t.Context(), task.tenant, task.session, task.turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err := journal.TransitionTurn(t.Context(), task.tenant, task.session, task.turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err == nil { err = journal.AppendTurnEvents(t.Context(), task.tenant, task.session, task.turn, 1, []sessions.ExecutionEvent{{Kind: "delta", Payload: json.RawMessage(`{"delta":"accepted"}`)}}) } diff --git a/services/core/internal/store/export_test.go b/services/core/internal/store/export_test.go index 89ad806ba..9be7be169 100644 --- a/services/core/internal/store/export_test.go +++ b/services/core/internal/store/export_test.go @@ -22,6 +22,15 @@ import ( func NewTestStore(t *testing.T) (*Store, *pgxpool.Pool) { return testStore(t) } +// SessionAdapter is the Session adapter on s's pool and credential key. +func SessionAdapter(s *Store) *sessionpg.Store { return sessionAdapter(s) } + +// TransitionTurn moves the Turn as the execution owner does, in a Session +// transaction on s's writer. +func TransitionTurn(ctx context.Context, s *Store, tenant, session, turn string, transition sessions.TurnTransition) (sessions.Turn, error) { + return transitionTurn(ctx, s, tenant, session, turn, transition) +} + var fixtureCipher, _ = credentialcrypto.New(bytes.Repeat([]byte{61}, 32)) // FixtureCipher is the credential key of NewModelTestStore, for reopened stores. diff --git a/services/core/internal/store/function_execution_native_test.go b/services/core/internal/store/function_execution_native_test.go index 58458ac01..a6ed06c0a 100644 --- a/services/core/internal/store/function_execution_native_test.go +++ b/services/core/internal/store/function_execution_native_test.go @@ -66,7 +66,7 @@ func TestNativeFunctionExecutionPersistsCallsResultsAndContinuity(t *testing.T) if !found { t.Fatal("required action identity differs from recovered function item") } - bound, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || bound.NativeSessionID == "" || (nativeID != "" && bound.NativeSessionID != nativeID) { t.Fatal(bound, err) } diff --git a/services/core/internal/store/function_execution_test.go b/services/core/internal/store/function_execution_test.go index b35aec657..c59f5f3a7 100644 --- a/services/core/internal/store/function_execution_test.go +++ b/services/core/internal/store/function_execution_test.go @@ -155,7 +155,7 @@ func TestExecutionFunctionsCancellationAndUnconfirmedResults(t *testing.T) { } h.finished(result, status) if cancel { - bound, err := h.s.GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) if err != nil || bound.NativeSessionID != "native-cancelled-functions" { t.Fatal(bound, err) } @@ -203,7 +203,7 @@ func TestExecutionFunctionsRequireAdvertisedCapability(t *testing.T) { if result.err == nil || !strings.Contains(result.err.Error(), "function_tools") { t.Fatal(result) } - turn, err := h.s.GetTurn(t.Context(), h.tenant, session.ID, input.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnQueued { t.Fatal(turn, err) } diff --git a/services/core/internal/store/function_images_native_test.go b/services/core/internal/store/function_images_native_test.go index 4d660edc6..bedafe512 100644 --- a/services/core/internal/store/function_images_native_test.go +++ b/services/core/internal/store/function_images_native_test.go @@ -82,14 +82,14 @@ func TestNativeFunctionImagePublicExecution(t *testing.T) { t.Fatal("function result admission duplicated or mutated", err) } } - before, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + before, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID == "" { t.Fatal("native binding missing", err) } stop() stop = startNativeEngineDaemon(t, h, home, binary, kind) run("resume") - after, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + after, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID != after.NativeSessionID { t.Fatal("native history changed", err) } diff --git a/services/core/internal/store/function_inputs_public_test.go b/services/core/internal/store/function_inputs_public_test.go index 72614c5d8..7570cbe0b 100644 --- a/services/core/internal/store/function_inputs_public_test.go +++ b/services/core/internal/store/function_inputs_public_test.go @@ -38,7 +38,7 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err := store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } functions := executionOwner(t, db, s).Sessions @@ -108,7 +108,7 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { if err != nil || len(history) != 6 { t.Fatal(history, err) } - if _, err := s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnWaiting, Status: sessions.TurnFailed}); err != nil { + if _, err := store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnWaiting, Status: sessions.TurnFailed}); err != nil { t.Fatal(err) } next, err := s.SubmitMessage(ctx, tenant, session.ID, "next", json.RawMessage(`{"text":"next"}`)) @@ -125,7 +125,7 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { if err != nil || len(history) != 1 { t.Fatal(history, err) } - current, err := s.GetTurn(ctx, tenant, session.ID, next.TurnID) + current, err := store.SessionAdapter(s).GetTurn(ctx, tenant, session.ID, next.TurnID) if err != nil || current.Status != sessions.TurnQueued || !current.CancelRequestedAt.IsZero() { t.Fatal(current, err) } diff --git a/services/core/internal/store/function_inputs_test.go b/services/core/internal/store/function_inputs_test.go index 53708dc1c..05384e8ab 100644 --- a/services/core/internal/store/function_inputs_test.go +++ b/services/core/internal/store/function_inputs_test.go @@ -81,7 +81,7 @@ func TestFunctionInputBatchesPersistAndReplayWithoutRetargeting(t *testing.T) { if err != nil || len(future) != 1 { t.Fatal(future, err) } - current, err := s.GetTurn(t.Context(), tenant, session.ID, next) + current, err := sessionAdapter(s).GetTurn(t.Context(), tenant, session.ID, next) if err != nil || !current.CancelRequestedAt.IsZero() || current.Status != sessions.TurnQueued { t.Fatal(current, err) } @@ -143,7 +143,7 @@ func TestFunctionInputBatchFailureRollsBackEveryWrite(t *testing.T) { if err != nil || call.Result != nil || call.Applied { t.Fatal(call, err) } - state, err := s.GetTurn(t.Context(), tenant, session.ID, turn) + state, err := sessionAdapter(s).GetTurn(t.Context(), tenant, session.ID, turn) if err != nil || !state.CancelRequestedAt.IsZero() || state.Status != sessions.TurnWaiting { t.Fatal(state, err) } diff --git a/services/core/internal/store/function_public_native_test.go b/services/core/internal/store/function_public_native_test.go index c96d5a1f8..1df47e840 100644 --- a/services/core/internal/store/function_public_native_test.go +++ b/services/core/internal/store/function_public_native_test.go @@ -49,7 +49,7 @@ func TestNativePublicFunctionExecution(t *testing.T) { t.Fatal(call, err) } } - bound, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || bound.NativeSessionID == "" || bound.Device.ID != h.device.ID { t.Fatal(bound, err) } diff --git a/services/core/internal/store/function_state_public_test.go b/services/core/internal/store/function_state_public_test.go index 96b8f8031..c578c5310 100644 --- a/services/core/internal/store/function_state_public_test.go +++ b/services/core/internal/store/function_state_public_test.go @@ -35,7 +35,7 @@ func TestFunctionStateOfficialClientReadsAndLiveEvents(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err := store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } functions := executionOwner(t, db, s).Sessions @@ -79,7 +79,7 @@ func TestFunctionStateOfficialClientReadsAndLiveEvents(t *testing.T) { t.Fatal(err) } } - if _, err := s.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCompleted, nil, "", input.Sequence); err != nil { + if _, err := functions.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCompleted, nil, "", input.Sequence); err != nil { t.Fatal(err) } if err := command.Wait(); err != nil { diff --git a/services/core/internal/store/function_state_test.go b/services/core/internal/store/function_state_test.go index 70dd61af6..24da5e4c5 100644 --- a/services/core/internal/store/function_state_test.go +++ b/services/core/internal/store/function_state_test.go @@ -38,7 +38,7 @@ func TestFunctionStateSnapshotsRecoveryAndRetries(t *testing.T) { pool.Close() s, _ = testStore(t) assertFunctionState(t, s, tenant, session.ID, sessions.TurnWaiting, 2) - if _, err := s.CompleteExecution(t.Context(), tenant, session.ID, turn, sessions.TurnCompleted, nil, "", 1); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := functions.CompleteExecution(t.Context(), tenant, session.ID, turn, sessions.TurnCompleted, nil, "", 1); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("waiting execution completed", err) } for i, id := range []string{"first", "second"} { @@ -110,7 +110,7 @@ func TestFunctionStateCancellationAndTerminalCleanup(t *testing.T) { } if status == sessions.TurnCompleted { transition(t, s, tenant, session.ID, input.TurnID, sessions.TurnWaiting, status) - } else if _, err := s.CompleteExecution(t.Context(), tenant, session.ID, input.TurnID, status, nil, "", input.Sequence); err != nil { + } else if _, err := functions.CompleteExecution(t.Context(), tenant, session.ID, input.TurnID, status, nil, "", input.Sequence); err != nil { t.Fatal(err) } assertFunctionState(t, s, tenant, session.ID, status, 0) diff --git a/services/core/internal/store/function_stream_native_test.go b/services/core/internal/store/function_stream_native_test.go index 9ddf2316e..11197f2f8 100644 --- a/services/core/internal/store/function_stream_native_test.go +++ b/services/core/internal/store/function_stream_native_test.go @@ -44,7 +44,7 @@ func TestNativePublicFunctionStreamHelper(t *testing.T) { t.Fatal(call, err) } } - bound, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || bound.NativeSessionID == "" || bound.Device.ID != h.device.ID { t.Fatal(bound, err) } diff --git a/services/core/internal/store/function_worker_test.go b/services/core/internal/store/function_worker_test.go index 00cfb8382..0d5efa0aa 100644 --- a/services/core/internal/store/function_worker_test.go +++ b/services/core/internal/store/function_worker_test.go @@ -79,7 +79,7 @@ func TestWorkerWaitsForToolCapabilities(t *testing.T) { } }() time.Sleep(650 * time.Millisecond) - current, err := h.s.GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) + current, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) if err != nil || current.Status != sessions.TurnQueued { t.Fatal(current, err) } diff --git a/services/core/internal/store/harness_onboarding_test.go b/services/core/internal/store/harness_onboarding_test.go index 69440c8dd..60879b0ff 100644 --- a/services/core/internal/store/harness_onboarding_test.go +++ b/services/core/internal/store/harness_onboarding_test.go @@ -22,6 +22,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -99,7 +100,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { } request("POST", "/v1/agents/sessions/"+created.ID+"/events", `{"events":[{"type":"agent.session.input.message","input":[{"role":"user","content":[{"type":"input_text","text":"finish"}]}]}]}`, 202) waitTurn(t, h, first.RunID, sessions.TurnCompleted) - turn, err := h.s.GetTurn(ctx, h.tenant, created.ID, first.RunID) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, created.ID, first.RunID) if err != nil { t.Fatal(err) } @@ -111,7 +112,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { if err = json.Unmarshal(turn.Outcome, &result); err != nil || result.AppliedThrough != inputs[1].Sequence || result.Done.Content != "readyfinish" { t.Fatal(string(turn.Outcome), err) } - bound, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, created.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, created.ID) if err != nil || bound.NativeSessionID == "" { t.Fatal(bound, err) } diff --git a/services/core/internal/store/input_batches_test.go b/services/core/internal/store/input_batches_test.go index 222816197..16574d0fd 100644 --- a/services/core/internal/store/input_batches_test.go +++ b/services/core/internal/store/input_batches_test.go @@ -108,7 +108,7 @@ func TestBatchRetriesCompareTheWholeRequestAndRetainTargets(t *testing.T) { if first[0].TurnID != "" || first[1].TurnID == "" || first[1].TurnID != first[2].TurnID || first[1].TurnID == first[3].TurnID { t.Fatalf("cancellation targets: %+v", first) } - cancelled, err := s.GetTurn(ctx, tenant, session.ID, first[1].TurnID) + cancelled, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, first[1].TurnID) if err != nil || cancelled.Status != sessions.TurnCancelled { t.Fatalf("cancelled turn=%+v err=%v", cancelled, err) } @@ -130,7 +130,7 @@ func TestBatchRetriesCompareTheWholeRequestAndRetainTargets(t *testing.T) { if err != nil || !reflect.DeepEqual(retry, first) { t.Fatalf("restart changed receipts: %+v, %v", retry, err) } - current, err := restarted.GetTurn(ctx, tenant, session.ID, next.TurnID) + current, err := sessionAdapter(restarted).GetTurn(ctx, tenant, session.ID, next.TurnID) if err != nil || current.Status != sessions.TurnQueued || !current.CancelRequestedAt.IsZero() { t.Fatalf("retry cancelled later work: %+v, %v", current, err) } @@ -158,7 +158,7 @@ func TestFailedBatchRollsBackEarlierCancellationAndInputs(t *testing.T) { if got, err := s.SubmitInputs(ctx, tenant, session.ID, key, batch); err == nil || got != nil { t.Fatalf("partial batch succeeded: %+v %v", got, err) } - turn, err := s.GetTurn(ctx, tenant, session.ID, initial.TurnID) + turn, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, initial.TurnID) if err != nil || turn.Status != sessions.TurnInProgress || !turn.CancelRequestedAt.IsZero() { t.Fatalf("cancellation escaped rollback: %+v %v", turn, err) } diff --git a/services/core/internal/store/input_conflicts_public_test.go b/services/core/internal/store/input_conflicts_public_test.go index fd1343bfb..fbe1e422c 100644 --- a/services/core/internal/store/input_conflicts_public_test.go +++ b/services/core/internal/store/input_conflicts_public_test.go @@ -74,7 +74,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.TransitionTurn(ctx, tenant, session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err := store.TransitionTurn(ctx, s, tenant, session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } if err := functions.RecordFunctionCall(ctx, tenant, session, receipt.TurnID, sessions.FunctionCall{CallID: call, ExecutorCallID: "native-" + call, Name: "lookup", Arguments: json.RawMessage(`{}`)}); err != nil { @@ -87,7 +87,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { if err := functions.ConfirmFunctionResult(ctx, tenant, session, receipt.TurnID, call); err != nil { t.Fatal(err) } - if _, err := s.CompleteExecution(ctx, tenant, session, receipt.TurnID, sessions.TurnCompleted, nil, "", receipt.Sequence); err != nil { + if _, err := functions.CompleteExecution(ctx, tenant, session, receipt.TurnID, sessions.TurnCompleted, nil, "", receipt.Sequence); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/item_order_migration_test.go b/services/core/internal/store/item_order_migration_test.go index 2d5691eef..dff7eb189 100644 --- a/services/core/internal/store/item_order_migration_test.go +++ b/services/core/internal/store/item_order_migration_test.go @@ -150,7 +150,7 @@ func TestItemOrderMigrationPreservesIndexedHistory(t *testing.T) { } t.Cleanup(migratedPool.Close) s := store.New(migratedPool) - if _, err = s.TransitionTurn(ctx, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err = store.TransitionTurn(ctx, s, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } if err = executionOwner(t, fixtureDB{pool: migratedPool}, s).Sessions.AppendTurnEvents(ctx, tenant, session, turn, 1, []sessions.ExecutionEvent{{Kind: "delta", Payload: json.RawMessage(`{"item_id":"after-upgrade","delta":"continued"}`)}}); err != nil { diff --git a/services/core/internal/store/item_order_test.go b/services/core/internal/store/item_order_test.go index 773d36aae..924da4b50 100644 --- a/services/core/internal/store/item_order_test.go +++ b/services/core/internal/store/item_order_test.go @@ -27,7 +27,7 @@ func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) if err != nil { t.Fatal(err) } - _, err = s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } @@ -128,7 +128,7 @@ func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) t.Fatalf("output index = %d, want %d", output.Int32, index) } } - if _, err = s.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{}`), "", input.Sequence); err != nil { + if _, err = journal.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{}`), "", input.Sequence); err != nil { t.Fatal(err) } checkOrder() @@ -136,7 +136,7 @@ func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) if err != nil { t.Fatal(err) } - if _, err = s.TransitionTurn(ctx, tenant, session.ID, next.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err = store.TransitionTurn(ctx, s, tenant, session.ID, next.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } if err = journal.AppendTurnEvents(ctx, tenant, session.ID, next.TurnID, 1, []sessions.ExecutionEvent{{Kind: "delta", Payload: json.RawMessage(`{"item_id":"new","delta":"new turn"}`)}}); err != nil { diff --git a/services/core/internal/store/item_reads_test.go b/services/core/internal/store/item_reads_test.go index f0f877c6b..836fd4333 100644 --- a/services/core/internal/store/item_reads_test.go +++ b/services/core/internal/store/item_reads_test.go @@ -27,7 +27,7 @@ func TestItemsRecoverSnapshotsPartialResultsPaginationAndIsolation(t *testing.T) if err != nil { t.Fatal(err) } - _, err = s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } @@ -57,7 +57,7 @@ func TestItemsRecoverSnapshotsPartialResultsPaginationAndIsolation(t *testing.T) if page.Items[4].Status != "in_progress" || page.Items[5].Status != "in_progress" { t.Fatal(page.Items) } - _, err = s.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{}`), "", input.Sequence) + _, err = journal.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{}`), "", input.Sequence) if err != nil { t.Fatal(err) } @@ -127,7 +127,7 @@ func TestItemProjectionFailureRollsBackJournalAndAggregateRecovers(t *testing.T) if err != nil { t.Fatal(err) } - _, err = s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } @@ -143,7 +143,7 @@ func TestItemProjectionFailureRollsBackJournalAndAggregateRecovers(t *testing.T) if err != nil || len(page.Items) != 1 { t.Fatal(page, err) } - _, err = s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnCompleted, Outcome: json.RawMessage(`{"done":{"content":"legacy answer","metadata":{"private":"SECRET"}}}`)}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnCompleted, Outcome: json.RawMessage(`{"done":{"content":"legacy answer","metadata":{"private":"SECRET"}}}`)}) if err != nil { t.Fatal(err) } @@ -164,7 +164,7 @@ func TestReceiptOnlyTextRecoversWithoutInventingCompletion(t *testing.T) { if err != nil { t.Fatal(err) } - _, err = s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } @@ -174,7 +174,7 @@ func TestReceiptOnlyTextRecoversWithoutInventingCompletion(t *testing.T) { t.Fatal(err) } } - _, err = s.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{"done":{"content":"retained cancellation text"}}`), "", input.Sequence) + _, err = journal.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{"done":{"content":"retained cancellation text"}}`), "", input.Sequence) if err != nil { t.Fatal(err) } @@ -198,7 +198,7 @@ func TestLegacyFailureRetainsPartialAnswerAcrossRecovery(t *testing.T) { if err != nil { t.Fatal(err) } - _, err = s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } @@ -212,7 +212,7 @@ func TestLegacyFailureRetainsPartialAnswerAcrossRecovery(t *testing.T) { if err = journal.AppendTurnEvents(ctx, tenant, session.ID, input.TurnID, 1, batch); err != nil { t.Fatal(err) } - _, err = s.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnFailed, json.RawMessage(`{"done":{"content":"provider failure"},"error_code":"engine_failed"}`), "", input.Sequence) + _, err = journal.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnFailed, json.RawMessage(`{"done":{"content":"provider failure"},"error_code":"engine_failed"}`), "", input.Sequence) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/list_cursor_public_test.go b/services/core/internal/store/list_cursor_public_test.go index 009578904..70d58d2dc 100644 --- a/services/core/internal/store/list_cursor_public_test.go +++ b/services/core/internal/store/list_cursor_public_test.go @@ -66,7 +66,7 @@ func seedCursorFixture(t *testing.T, s *store.Store, leased execution.Owner, ski f.session = client.created(token, "/v1/agents/sessions", newSession) f.turn = first("/v1/agents/sessions/" + f.session + "/turns") f.item = first("/v1/agents/sessions/" + f.session + "/items") - if _, err := s.TransitionTurn(ctx, tenant, f.session, f.turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnCancelled}); err != nil { + if _, err := store.TransitionTurn(ctx, s, tenant, f.session, f.turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnCancelled}); err != nil { t.Fatal(err) } if _, err := s.SubmitMessage(ctx, tenant, f.session, label+"-second", json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"second"}]}]}`)); err != nil { @@ -102,7 +102,7 @@ func seedCursorFixture(t *testing.T, s *store.Store, leased execution.Owner, ski if err = leased.Sessions.BindSessionDevice(ctx, tenant, created.ID, host.ID); err != nil { t.Fatal(err) } - if _, err = leased.Store.TransitionTurn(ctx, tenant, created.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err = store.TransitionTurn(ctx, leased.Store, tenant, created.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } opened := int64(1700000001000) diff --git a/services/core/internal/store/local_environment_devices_test.go b/services/core/internal/store/local_environment_devices_test.go index 00c31eff8..a16b46b3d 100644 --- a/services/core/internal/store/local_environment_devices_test.go +++ b/services/core/internal/store/local_environment_devices_test.go @@ -47,7 +47,7 @@ func TestEnvironmentDeviceAuthorityAndLifecycle(t *testing.T) { t.Fatal("dedicated credential bound to another Session") } } - devices, err := s.ListExecutionDevices(t.Context(), tenant) + devices, err := sessionAdapter(s).ListExecutionDevices(t.Context(), tenant) if err != nil || len(devices) != 0 { t.Fatalf("dedicated device entered general selection: %v %v", devices, err) } diff --git a/services/core/internal/store/local_environment_worker_test.go b/services/core/internal/store/local_environment_worker_test.go index 1aabca710..001bec8cb 100644 --- a/services/core/internal/store/local_environment_worker_test.go +++ b/services/core/internal/store/local_environment_worker_test.go @@ -141,14 +141,14 @@ func TestLocalEnvironmentWorkerSchedulesPreparationWithoutRemoteResolver(t *test h.write(start.RunID, proto.TypeDone, proto.DonePayload{Content: "complete", Metadata: map[string]any{proto.DoneMetaAgentSessionID: "local-native-history"}}) completeLocalArtifactExport(t, h, worker, environment) awaitDaemonRemoteCondition(t, t.Context(), 5*time.Second, "local completion", func() bool { - turn, err := h.s.GetTurn(t.Context(), h.tenant, h.session.ID, start.RunID) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, start.RunID) return err == nil && turn.Status == sessions.TurnCompleted }) settled, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, reservation.ID) if err != nil || settled.State != sessions.EnvironmentInputAdmitted || len(settled.Receipts) != 1 { t.Fatal("local reservation did not settle", err) } - bound, err := h.s.GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) if err != nil || bound.Device.EnvironmentID != environment.ID || bound.NativeSessionID != "local-native-history" { t.Fatal("local native identity was not retained", err) } diff --git a/services/core/internal/store/mcode_public_native_test.go b/services/core/internal/store/mcode_public_native_test.go index 036eab079..2ffb143e3 100644 --- a/services/core/internal/store/mcode_public_native_test.go +++ b/services/core/internal/store/mcode_public_native_test.go @@ -13,6 +13,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -64,7 +65,7 @@ func TestNativeMCodePublicExecution(t *testing.T) { Session string `json:"session"` } _ = json.Unmarshal(data, &identity) - if page, e := h.s.ListTurns(ctx, h.tenant, identity.Session, "", 100, true); e == nil { + if page, e := store.SessionAdapter(h.s).ListTurns(ctx, h.tenant, identity.Session, "", 100, true); e == nil { diagnostic, _ := json.Marshal(page) text := strings.ReplaceAll(string(diagnostic), provider.APIKey, "[REDACTED]") _ = os.WriteFile(filepath.Join(home, "failed-turns.json"), []byte(text), 0600) @@ -84,7 +85,7 @@ func TestNativeMCodePublicExecution(t *testing.T) { if json.Unmarshal(data, &proof) != nil { t.Fatal("invalid evidence") } - turn, err := h.s.GetTurn(ctx, h.tenant, proof.Session, proof.FirstTurn) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, proof.Session, proof.FirstTurn) if err != nil { t.Fatal(err) } @@ -96,14 +97,14 @@ func TestNativeMCodePublicExecution(t *testing.T) { if json.Unmarshal(turn.Outcome, &outcome) != nil || outcome.AppliedThrough != inputs[1].Sequence { t.Fatal("native applied receipt missing") } - before, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + before, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID == "" { t.Fatal("native binding missing", err) } stop() stop = startNativeEngineDaemon(t, h, home, binary, "mcode") run("resume") - after, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + after, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID != after.NativeSessionID { t.Fatal("native history changed", err) } diff --git a/services/core/internal/store/mcp_credential_selection_public_test.go b/services/core/internal/store/mcp_credential_selection_public_test.go index 88837bdb2..1915c220d 100644 --- a/services/core/internal/store/mcp_credential_selection_public_test.go +++ b/services/core/internal/store/mcp_credential_selection_public_test.go @@ -282,7 +282,7 @@ func TestMCPCredentialSelectionPublicPostgres(t *testing.T) { if err := pool.QueryRow(t.Context(), "SELECT id FROM turns WHERE session_id=$1", streamed).Scan(&turn); err != nil { t.Fatal(err) } - if _, err := s.TransitionTurn(t.Context(), tenantA, streamed, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnCancelled}); err != nil { + if _, err := store.TransitionTurn(t.Context(), s, tenantA, streamed, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnCancelled}); err != nil { t.Fatal(err) } _, tools = snapshot(stream, "agent.session.idle") diff --git a/services/core/internal/store/message_images_native_test.go b/services/core/internal/store/message_images_native_test.go index f59e3b23b..46c2583cb 100644 --- a/services/core/internal/store/message_images_native_test.go +++ b/services/core/internal/store/message_images_native_test.go @@ -78,7 +78,7 @@ func TestNativeMessageImagePublicExecution(t *testing.T) { if err != nil || !call.Applied { t.Fatal("function application receipt missing", err) } - turn, err := h.s.GetTurn(ctx, h.tenant, proof.Session, proof.Turn) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, proof.Session, proof.Turn) if err != nil { t.Fatal(err) } @@ -90,14 +90,14 @@ func TestNativeMessageImagePublicExecution(t *testing.T) { if err != nil || len(inputs) != 3 || inputs[0].Kind != "message" || inputs[1].Kind != "message" || inputs[2].Kind != "tool_result" || outcome.AppliedThrough != inputs[2].Sequence { t.Fatal("active image batch was not applied exactly once in the same Turn", err) } - before, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + before, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID == "" { t.Fatal("native binding missing", err) } stop() stop = startNativeEngineDaemon(t, h, home, binary, kind) run("resume") - after, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + after, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID != after.NativeSessionID { t.Fatal("native history changed", err) } diff --git a/services/core/internal/store/model_protocol_native_test.go b/services/core/internal/store/model_protocol_native_test.go index c46ab363d..dceecd9bf 100644 --- a/services/core/internal/store/model_protocol_native_test.go +++ b/services/core/internal/store/model_protocol_native_test.go @@ -136,14 +136,14 @@ func TestNativeModelProtocolPublicExecution(t *testing.T) { if expectedCalls > 0 && failed != 1 { t.Fatal("missing failed function-result scenario") } - before, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + before, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID == "" { t.Fatal("native Session binding missing before restart") } stop() stop = startNativeEngineDaemon(t, h, home, binary, options.Engine) run("resume") - after, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + after, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID != after.NativeSessionID { t.Fatal("cold Session continuation changed native history") } diff --git a/services/core/internal/store/native_environment_test.go b/services/core/internal/store/native_environment_test.go index e86f74c55..547ad5f2c 100644 --- a/services/core/internal/store/native_environment_test.go +++ b/services/core/internal/store/native_environment_test.go @@ -108,13 +108,13 @@ func TestNativeNoExecutionEnvironment(t *testing.T) { } first := h.message("first", "Return an answer.") h.finished(h.run(ctx, first.TurnID), sessions.TurnCompleted) - bound, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || bound.NativeSessionID == "" { t.Fatal(bound, err) } second := h.message("second", "Continue the same conversation.") h.finished(h.run(ctx, second.TurnID), sessions.TurnCompleted) - again, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) + again, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || again.NativeSessionID != bound.NativeSessionID { t.Fatal(again, err) } diff --git a/services/core/internal/store/native_recovery_test.go b/services/core/internal/store/native_recovery_test.go index f95a1293e..7cf357098 100644 --- a/services/core/internal/store/native_recovery_test.go +++ b/services/core/internal/store/native_recovery_test.go @@ -23,11 +23,11 @@ func TestSessionExecutionBindingRetainsStartedExecutionRequirement(t *testing.T) } assertStarted := func(st *Store, want bool) { t.Helper() - bound, err := st.GetSessionExecutionBinding(t.Context(), tenant, session.ID) + bound, err := sessionAdapter(st).GetSessionExecutionBinding(t.Context(), tenant, session.ID) if err != nil || bound.HasStartedTurn != want || bound.NativeSessionID != "" { t.Fatalf("binding=%+v err=%v", bound, err) } - if _, err := st.GetSessionExecutionBinding(t.Context(), foreign, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(st).GetSessionExecutionBinding(t.Context(), foreign, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign binding", err) } } diff --git a/services/core/internal/store/no_environment_test.go b/services/core/internal/store/no_environment_test.go index c43cf3fed..6c32bf2c1 100644 --- a/services/core/internal/store/no_environment_test.go +++ b/services/core/internal/store/no_environment_test.go @@ -23,7 +23,7 @@ func TestNoEnvironmentRejectsUnadvertisedDeviceBeforeClaim(t *testing.T) { if _, err = h.bound().Run(ctx, h.tenant, h.session.ID, input.TurnID); err == nil { t.Fatal("unsupported environment admitted") } - turn, err := h.s.GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnQueued { t.Fatal(turn, err) } diff --git a/services/core/internal/store/prepared_dispatch_failure_test.go b/services/core/internal/store/prepared_dispatch_failure_test.go index 5ed6a8480..23f05ff48 100644 --- a/services/core/internal/store/prepared_dispatch_failure_test.go +++ b/services/core/internal/store/prepared_dispatch_failure_test.go @@ -190,7 +190,7 @@ func TestPreparedDispatchCancellationReceiptSurvivesStartFailure(t *testing.T) { if receipts != 1 { t.Fatal("cancellation receipt was not journaled once", receipts) } - bound, err := h.s.GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) if err != nil || (withOutcome && (bound.NativeSessionID != "cancelled-prepared-native" || outcome.Done.Content != "retained cancellation")) { t.Fatal("cancellation lost native continuation or final output", err) } diff --git a/services/core/internal/store/prepared_dispatch_test.go b/services/core/internal/store/prepared_dispatch_test.go index a05dac21e..109b662bf 100644 --- a/services/core/internal/store/prepared_dispatch_test.go +++ b/services/core/internal/store/prepared_dispatch_test.go @@ -66,7 +66,7 @@ func readyPreparedDispatch(t *testing.T, h *dispatchHarness, request, handle str if frame.ID != request || frame.DecodePayload(&start) != nil || start.Handle != handle || start.RunID == "" || inputTextForTest(t, start.Input) != "first\n\nsecond" { t.Fatal("Start changed preparation or original batch", frame.ID, start) } - turn, err := h.s.GetTurn(t.Context(), h.tenant, h.session.ID, start.RunID) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, start.RunID) if err != nil || turn.Status != sessions.TurnInProgress { t.Fatal("Start preceded atomic claim", turn, err) } @@ -106,7 +106,7 @@ func TestPreparedDispatchPromotesOriginalBatchAndPersistsCompletion(t *testing.T t.Fatal("prepared completion", got) } assertPreparationReleased(t, h, frame.ID, handle) - bound, err := h.s.GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(t.Context(), h.tenant, h.session.ID) if err != nil || bound.NativeSessionID != "retained-prepared-native" { t.Fatal("native identity was not committed", bound, err) } diff --git a/services/core/internal/store/public_execution_test.go b/services/core/internal/store/public_execution_test.go index 9c8df0da2..8d3301506 100644 --- a/services/core/internal/store/public_execution_test.go +++ b/services/core/internal/store/public_execution_test.go @@ -95,7 +95,7 @@ func waitTurn(t *testing.T, h *dispatchHarness, id, status string) { t.Helper() deadline := time.Now().Add(12 * time.Second) for time.Now().Before(deadline) { - turn, err := h.s.GetTurn(context.Background(), h.tenant, h.session.ID, id) + turn, err := store.SessionAdapter(h.s).GetTurn(context.Background(), h.tenant, h.session.ID, id) if err != nil { t.Fatal(err) } @@ -112,7 +112,7 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { h.session = publicSession(t, h, "interrupted") first := h.message("first", "Already sent") ctx := context.Background() - if _, err := h.s.TransitionTurn(ctx, h.tenant, h.session.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err := store.TransitionTurn(ctx, h.s, h.tenant, h.session.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } // A native measurement committed before process loss must survive startup @@ -123,7 +123,7 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { } checkMeasurement := func(ended bool) { t.Helper() - turn, err := h.s.GetTurn(ctx, h.tenant, h.session.ID, first.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, h.session.ID, first.TurnID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/public_handler_fixture_test.go b/services/core/internal/store/public_handler_fixture_test.go index b04fa9f73..111eeb790 100644 --- a/services/core/internal/store/public_handler_fixture_test.go +++ b/services/core/internal/store/public_handler_fixture_test.go @@ -88,7 +88,7 @@ func publicHandler(t testing.TB, s *store.Store, db fixtureDB, keys fixtureKeyRe Sessions: s, SessionCreation: s, SessionEvents: s, - Turns: s, + Turns: sessionStore, Items: sessionStore, Subagents: sessionStore, Artifacts: sessionService, diff --git a/services/core/internal/store/runtime_capabilities_pending_test.go b/services/core/internal/store/runtime_capabilities_pending_test.go index 1052881c3..06cb25c34 100644 --- a/services/core/internal/store/runtime_capabilities_pending_test.go +++ b/services/core/internal/store/runtime_capabilities_pending_test.go @@ -52,7 +52,7 @@ func TestManagedCapabilitiesWaitBeforeInitializationClaim(t *testing.T) { if err != nil || initializationState(t, db.pool, owner.TenantID, owner.EnvironmentID) != "pending" || owner.State != "running" || provider.writes.Load() != 0 || provider.kills != 0 { t.Fatal("missing socket consumed initialization or requested cleanup", owner, err, provider.writes.Load(), provider.kills) } - if _, err := s.GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := store.SessionAdapter(s).GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("ordinary readiness gate bypassed", err) } } diff --git a/services/core/internal/store/runtime_enrollment_test.go b/services/core/internal/store/runtime_enrollment_test.go index 2a4a001ba..ab2e10932 100644 --- a/services/core/internal/store/runtime_enrollment_test.go +++ b/services/core/internal/store/runtime_enrollment_test.go @@ -53,7 +53,7 @@ func TestRuntimeEnrollmentAuthorityAndRotation(t *testing.T) { if again, err := sessionService(t, s).EnrollRuntime(ctx, environment.ID, executorDigest(key.Token)); err != nil || again != bound { t.Fatalf("retry changed binding: %+v %v", again, err) } - if devices, err := s.ListExecutionDevices(ctx, p.TenantID); err != nil || len(devices) != 0 { + if devices, err := sessionAdapter(s).ListExecutionDevices(ctx, p.TenantID); err != nil || len(devices) != 0 { t.Fatalf("enrolled Runtime entered general selection: %v", err) } auth := runtimegateway.NewAuthenticator(sessionAdapter(s)) diff --git a/services/core/internal/store/runtime_idle_clock_test.go b/services/core/internal/store/runtime_idle_clock_test.go index 84e59a4b2..0cf23151e 100644 --- a/services/core/internal/store/runtime_idle_clock_test.go +++ b/services/core/internal/store/runtime_idle_clock_test.go @@ -87,7 +87,7 @@ func TestManagedIdleClockIgnoresRootHostSkew(t *testing.T) { source := runtimeDatabaseTime(t, s).Add(skew).UnixMilli() outcome := json.RawMessage(fmt.Sprintf(`{"done":{"source_completed_at_ms":%d}}`, source)) before := runtimeDatabaseTime(t, s) - completed, err := w.CompleteExecution(t.Context(), owner.TenantID, owner.SessionID, turn, sessions.TurnCompleted, outcome, "", 0) + completed, err := completeExecution(t.Context(), t, w, owner.TenantID, owner.SessionID, turn, sessions.TurnCompleted, outcome, "", 0) after := runtimeDatabaseTime(t, s) if err != nil || completed.CompletedAt.UnixMilli() != source { t.Fatal("native completion changed or rejected", completed, err) @@ -96,7 +96,7 @@ func TestManagedIdleClockIgnoresRootHostSkew(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := w.CompleteExecution(t.Context(), owner.TenantID, owner.SessionID, turn, sessions.TurnCompleted, outcome, "", 0); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := completeExecution(t.Context(), t, w, owner.TenantID, owner.SessionID, turn, sessions.TurnCompleted, outcome, "", 0); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("terminal replay accepted", err) } unchanged, err := deploymentStore(w).Activity(t.Context(), owner.ID) @@ -104,7 +104,7 @@ func TestManagedIdleClockIgnoresRootHostSkew(t *testing.T) { t.Fatal("terminal retry reset idle", unchanged, err) } verifyManagedIdleClock(t, s, w, owner, before, after) - read, err := s.GetTurn(t.Context(), owner.TenantID, owner.SessionID, turn) + read, err := sessionAdapter(s).GetTurn(t.Context(), owner.TenantID, owner.SessionID, turn) if err != nil || read.CompletedAt.UnixMilli() != source { t.Fatal("public native timestamp rewritten", read, err) } @@ -178,15 +178,15 @@ func TestUnmanagedRootCompletionPreservesHostSkew(t *testing.T) { current := transition(t, s, tenant, session.ID, input.TurnID, sessions.TurnQueued, sessions.TurnInProgress) source := current.CreatedAt.Add(skew).UnixMilli() outcome := json.RawMessage(fmt.Sprintf(`{"done":{"source_completed_at_ms":%d}}`, source)) - completed, err := s.CompleteExecution(t.Context(), tenant, session.ID, input.TurnID, sessions.TurnCompleted, outcome, "", input.Sequence) + completed, err := completeExecution(t.Context(), t, s, tenant, session.ID, input.TurnID, sessions.TurnCompleted, outcome, "", input.Sequence) if err != nil || completed.CompletedAt.UnixMilli() != source { t.Fatal("native completion changed or rejected", completed, err) } - read, err := s.GetTurn(t.Context(), tenant, session.ID, input.TurnID) + read, err := sessionAdapter(s).GetTurn(t.Context(), tenant, session.ID, input.TurnID) if err != nil || read.Status != sessions.TurnCompleted || read.CompletedAt.UnixMilli() != source { t.Fatal("public native timestamp rewritten", read, err) } - if _, err := s.CompleteExecution(t.Context(), tenant, session.ID, input.TurnID, sessions.TurnCompleted, outcome, "", input.Sequence); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := completeExecution(t.Context(), t, s, tenant, session.ID, input.TurnID, sessions.TurnCompleted, outcome, "", input.Sequence); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("terminal replay accepted", err) } }) @@ -202,7 +202,7 @@ func TestRootCompletionRejectsNonpositiveSourceTime(t *testing.T) { input := submitMessage(t, s, tenant, session.ID, "invalid-clock") transition(t, s, tenant, session.ID, input.TurnID, sessions.TurnQueued, sessions.TurnInProgress) outcome := json.RawMessage(fmt.Sprintf(`{"done":{"source_completed_at_ms":%d}}`, source)) - if _, err := s.CompleteExecution(t.Context(), tenant, session.ID, input.TurnID, sessions.TurnCompleted, outcome, "", input.Sequence); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := completeExecution(t.Context(), t, s, tenant, session.ID, input.TurnID, sessions.TurnCompleted, outcome, "", input.Sequence); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal("invalid native timestamp accepted", err) } }) @@ -221,7 +221,7 @@ func TestManagedIdleClockReconnectPreservesReceipts(t *testing.T) { if err != nil || !reflect.DeepEqual(before, after) { t.Fatal("reconnect changed compute receipt or idle clock", err) } - if _, err := s.GetTurn(t.Context(), owner.TenantID, owner.SessionID, turn); err != nil { + if _, err := sessionAdapter(s).GetTurn(t.Context(), owner.TenantID, owner.SessionID, turn); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/runtime_initialization.go b/services/core/internal/store/runtime_initialization.go deleted file mode 100644 index d156ef9ad..000000000 --- a/services/core/internal/store/runtime_initialization.go +++ /dev/null @@ -1,25 +0,0 @@ -package store - -import ( - "context" - "errors" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -func (s *Store) requireInitializedEnvironment(ctx context.Context, tenant, session pgtype.UUID) error { - ready, err := s.queries.GetSessionInitializationReady(ctx, sqlc.GetSessionInitializationReadyParams{TenantID: tenant, ID: session}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - if !ready { - return sessions.ErrNotFound - } - return nil -} diff --git a/services/core/internal/store/runtime_initialization_test.go b/services/core/internal/store/runtime_initialization_test.go index 7d61c1ebf..c50208edf 100644 --- a/services/core/internal/store/runtime_initialization_test.go +++ b/services/core/internal/store/runtime_initialization_test.go @@ -79,7 +79,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { if _, err := fixtureSessionStore(db).GetSessionDevice(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Error("premature file access", err) } - if _, err := s.GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := store.SessionAdapter(s).GetSessionExecutionBinding(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Error("premature execution", err) } if mode == "uncertain" { @@ -120,7 +120,7 @@ func TestEnvironmentInitializationCompletionUnknownAndRestart(t *testing.T) { if int(p.writes.Load()) != expectedSteps { t.Fatal("missing operations", p.writes.Load()) } - if _, err := s.GetSessionExecutionBinding(t.Context(), tenant, session.ID); err != nil { + if _, err := store.SessionAdapter(s).GetSessionExecutionBinding(t.Context(), tenant, session.ID); err != nil { t.Fatal("completed preparation blocked", err) } stop() diff --git a/services/core/internal/store/runtime_suspension_concurrency_test.go b/services/core/internal/store/runtime_suspension_concurrency_test.go index 5d2a97fd7..67aaf940f 100644 --- a/services/core/internal/store/runtime_suspension_concurrency_test.go +++ b/services/core/internal/store/runtime_suspension_concurrency_test.go @@ -61,7 +61,7 @@ func TestRuntimeSuspensionClaimReadsPhaseAfterSessionLock(t *testing.T) { ctx, tx, blocker := runtimeSuspensionLockedSession(t, pool, owner.SessionID) done := make(chan error, 1) go func() { - _, err := w.TransitionTurn(ctx, owner.TenantID, owner.SessionID, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err := transitionTurn(ctx, w, owner.TenantID, owner.SessionID, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) done <- err }() runtimeSuspensionWaitBlocked(t, ctx, pool, blocker, done) @@ -77,7 +77,7 @@ func TestRuntimeSuspensionClaimReadsPhaseAfterSessionLock(t *testing.T) { if blocked && !errors.Is(err, sessions.ErrTurnConflict) || !blocked && err != nil { t.Fatal("incorrect claim outcome", phase, err) } - got, err := s.GetTurn(ctx, owner.TenantID, owner.SessionID, turn) + got, err := sessionAdapter(s).GetTurn(ctx, owner.TenantID, owner.SessionID, turn) if err != nil { t.Fatal(err) } @@ -205,7 +205,7 @@ func TestRuntimeSuspensionQuiesceCannotOvertakeClaim(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := transitionTurn(ctx, w.queries.WithTx(tx), params, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress, Outcome: json.RawMessage(`{}`)}); err != nil { + if _, err := sessions.TransitionTurn(ctx, sessionpg.BindSession(w.queries.WithTx(tx), params.TenantID, params.SessionID), turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } if err := tx.Commit(ctx); err != nil { diff --git a/services/core/internal/store/runtime_worker_recovery_test.go b/services/core/internal/store/runtime_worker_recovery_test.go index db988287d..56853bb93 100644 --- a/services/core/internal/store/runtime_worker_recovery_test.go +++ b/services/core/internal/store/runtime_worker_recovery_test.go @@ -124,7 +124,7 @@ func TestWorkerWaitsForComputeAndSurvivesPromotionConflict(t *testing.T) { t.Fatal("completed preparation was not released") } awaitDaemonRemoteCondition(t, t.Context(), 5*time.Second, "original input completed once", func() bool { - turn, err := h.s.GetTurn(t.Context(), h.tenant, h.session.ID, start.RunID) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, start.RunID) return err == nil && turn.Status == sessions.TurnCompleted }) var turns int @@ -154,7 +154,7 @@ func TestWorkerRestartPreservesQueuedTurnWhileComputeWakes(t *testing.T) { if err := worker.Run(ctx); !errors.Is(err, context.Canceled) { t.Fatal(err) } - got, err := h.s.GetTurn(t.Context(), h.tenant, h.session.ID, turn) + got, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, turn) if err != nil || got.Status != sessions.TurnQueued || !got.StartedAt.IsZero() { t.Fatal("startup consumed queued work before restore", got, err) } diff --git a/services/core/internal/store/sandbox_deployment_switch_test.go b/services/core/internal/store/sandbox_deployment_switch_test.go index 3447ee77d..42848e08b 100644 --- a/services/core/internal/store/sandbox_deployment_switch_test.go +++ b/services/core/internal/store/sandbox_deployment_switch_test.go @@ -315,13 +315,13 @@ func TestSandboxSwitchPreservesReleasedAllocationAndItemHistory(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := w.TransitionTurn(t.Context(), tenant, history.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err := transitionTurn(t.Context(), w, tenant, history.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } if err := sessionExecution(t, w.lease).AppendTurnEvents(t.Context(), tenant, history.ID, input.TurnID, 1, []sessions.ExecutionEvent{{Kind: "output_message", Payload: json.RawMessage(`{"id":"answer","status":"completed","text":"retained answer"}`)}}); err != nil { t.Fatal(err) } - if _, err := w.TransitionTurn(t.Context(), tenant, history.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnCompleted}); err != nil { + if _, err := transitionTurn(t.Context(), w, tenant, history.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnCompleted}); err != nil { t.Fatal(err) } items, err := sessionAdapter(s).ListItems(t.Context(), tenant, history.ID, "", 100, true) diff --git a/services/core/internal/store/scheduling.go b/services/core/internal/store/scheduling.go index 0cb7a446b..c02516880 100644 --- a/services/core/internal/store/scheduling.go +++ b/services/core/internal/store/scheduling.go @@ -11,31 +11,10 @@ import ( "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 executionWorkCursor(after string, connectedDevices []string) (pgtype.UUID, []pgtype.UUID, error) { - id := pgtype.UUID{Valid: true} - var err error - if after != "" { - id, err = parseID(after) - if err != nil { - return id, nil, err - } - } - devices := make([]pgtype.UUID, 0, len(connectedDevices)) - for _, value := range connectedDevices { - device, err := parseID(value) - if err != nil { - return id, nil, err - } - devices = append(devices, device) - } - return id, devices, nil -} - func (s *Store) ListEnvironmentInputWork(ctx context.Context, after string, connectedDevices []string) ([]sessions.EnvironmentInputWork, error) { - id, devices, err := executionWorkCursor(after, connectedDevices) + id, devices, err := sessionpg.ExecutionWorkCursor(after, connectedDevices) if err != nil { return nil, err } @@ -50,38 +29,6 @@ func (s *Store) ListEnvironmentInputWork(ctx context.Context, after string, conn return work, nil } -func (s *Store) ListExecutionWork(ctx context.Context, after string, statuses []string, connectedDevices []string) ([]sessions.ExecutionWork, error) { - id, devices, err := executionWorkCursor(after, connectedDevices) - if err != nil { - return nil, err - } - rows, err := s.queries.ListExecutionWork(ctx, sqlc.ListExecutionWorkParams{AfterID: id, Statuses: statuses, ConnectedOnly: connectedDevices != nil, ConnectedDevices: devices}) - if err != nil { - return nil, err - } - work := make([]sessions.ExecutionWork, 0, len(rows)) - for _, row := range rows { - work = append(work, sessions.ExecutionWork{TenantID: uuid.UUID(row.TenantID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), TurnID: uuid.UUID(row.ID.Bytes).String(), Status: row.Status}) - } - return work, nil -} - -func (s *Store) ListExecutionDevices(ctx context.Context, tenantID string) ([]sessions.ExecutionDevice, error) { - tenant, err := parseID(tenantID) - if err != nil { - return nil, err - } - rows, err := s.queries.ListExecutionDevices(ctx, tenant) - if err != nil { - return nil, err - } - devices := make([]sessions.ExecutionDevice, 0, len(rows)) - for _, row := range rows { - devices = append(devices, sessions.ExecutionDevice{ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name}) - } - return devices, nil -} - func (s *Store) sessionActivity(ctx context.Context, session sessions.Session, err error) (sessions.Session, error) { if err != nil { return sessions.Session{}, err diff --git a/services/core/internal/store/self_hosted_cancel_public_test.go b/services/core/internal/store/self_hosted_cancel_public_test.go index 8a9469d86..eec981f6b 100644 --- a/services/core/internal/store/self_hosted_cancel_public_test.go +++ b/services/core/internal/store/self_hosted_cancel_public_test.go @@ -15,6 +15,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -124,7 +125,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { pending := map[string]string{created.InitialID: snapshot(created.InitialID), created.LaterID: snapshot(created.LaterID)} transition := func(id, from, to string) { t.Helper() - if _, err := s.TransitionTurn(t.Context(), tenant, created.ID, id, sessions.TurnTransition{ExpectedStatus: from, Status: to}); err != nil { + if _, err := store.TransitionTurn(t.Context(), s, tenant, created.ID, id, sessions.TurnTransition{ExpectedStatus: from, Status: to}); err != nil { t.Fatal(err) } } @@ -158,7 +159,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { } run("active") activeReceipts := receipts(created.ActiveKey, first) - turn, err := s.GetTurn(t.Context(), tenant, created.ID, first) + turn, err := store.SessionAdapter(s).GetTurn(t.Context(), tenant, created.ID, first) if err != nil || turn.Status != sessions.TurnInProgress || turn.CancelRequestedAt.IsZero() || !turn.CompletedAt.IsZero() { t.Fatal("202 must admit cancellation without fabricating native completion", err) } diff --git a/services/core/internal/store/session_artifacts_public_test.go b/services/core/internal/store/session_artifacts_public_test.go index 55c5f54eb..8447e86de 100644 --- a/services/core/internal/store/session_artifacts_public_test.go +++ b/services/core/internal/store/session_artifacts_public_test.go @@ -43,7 +43,7 @@ func completeArtifactTurn(t *testing.T, s *store.Store, artifacts *sessions.Serv } transition := func(from, to string) { t.Helper() - if _, err := s.TransitionTurn(t.Context(), tenant, session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: from, Status: to}); err != nil { + if _, err := store.TransitionTurn(t.Context(), s, tenant, session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: from, Status: to}); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/session_artifacts_test.go b/services/core/internal/store/session_artifacts_test.go index 2504d2916..0c4fd6d8d 100644 --- a/services/core/internal/store/session_artifacts_test.go +++ b/services/core/internal/store/session_artifacts_test.go @@ -449,7 +449,7 @@ func TestSessionArtifactsNewestVersionFollowsTurnOrder(t *testing.T) { transition(t, s, tenant, session, first.TurnID, sessions.TurnQueued, sessions.TurnInProgress) stageArtifactOutputs(t, s, tenant, session, env.ID, first.TurnID, map[string]string{"b.txt": "bravo"}) future := time.Now().Add(time.Hour).UnixMilli() - if _, err := s.CompleteExecution(t.Context(), tenant, session, first.TurnID, sessions.TurnCompleted, json.RawMessage(fmt.Sprintf(`{"done":{"source_completed_at_ms":%d}}`, future)), "", first.Sequence); err != nil { + if _, err := completeExecution(t.Context(), t, s, tenant, session, first.TurnID, sessions.TurnCompleted, json.RawMessage(fmt.Sprintf(`{"done":{"source_completed_at_ms":%d}}`, future)), "", first.Sequence); err != nil { t.Fatal(err) } one := publishedByTurn(t, s, tenant, session, first.TurnID)["b.txt"] @@ -510,7 +510,7 @@ func TestSessionArtifactsCompletionWaitsForConcurrentDeletion(t *testing.T) { } done := make(chan error, 1) go func() { - _, err := s.TransitionTurn(t.Context(), tenant, session, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnCompleted}) + _, err := transitionTurn(t.Context(), s, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnCompleted}) done <- err }() // Completion must be blocked on the Session lock before the deletion commits. @@ -533,7 +533,7 @@ func TestSessionArtifactsCompletionWaitsForConcurrentDeletion(t *testing.T) { } time.Sleep(10 * time.Millisecond) } - if status, err := s.GetTurn(t.Context(), tenant, session, turn); err != nil || status.Status != sessions.TurnInProgress { + if status, err := sessionAdapter(s).GetTurn(t.Context(), tenant, session, turn); err != nil || status.Status != sessions.TurnInProgress { t.Fatalf("Turn settled while the deletion held the lock: %+v %v", status, err) } if err := tx.Commit(t.Context()); err != nil { diff --git a/services/core/internal/store/session_creation_identity_test.go b/services/core/internal/store/session_creation_identity_test.go index b87bc54d1..28df45d2a 100644 --- a/services/core/internal/store/session_creation_identity_test.go +++ b/services/core/internal/store/session_creation_identity_test.go @@ -197,7 +197,7 @@ func TestSessionCreationKeepsItsResolvedDeploymentRevision(t *testing.T) { receipt := submitMessage(t, s, tenant, session, uuid.NewString()) transition(t, s, tenant, session, receipt.TurnID, sessions.TurnQueued, sessions.TurnInProgress) outcome, _ := json.Marshal(map[string]string{"error_code": coreCode, "engine_error_code": nativeCode}) - turn, err := s.CompleteExecution(t.Context(), tenant, session, receipt.TurnID, status, outcome, "", receipt.Sequence) + turn, err := completeExecution(t.Context(), t, s, tenant, session, receipt.TurnID, status, outcome, "", receipt.Sequence) if err != nil { t.Fatal(err) } @@ -233,7 +233,7 @@ func TestSessionCreationKeepsItsResolvedDeploymentRevision(t *testing.T) { if revision(stale.ID) != original { t.Fatal("tuple revision changed during insertion") } - frozen, err := s.SessionModelExecution(t.Context(), tenant, stale.ID) + frozen, err := sessionAdapter(s).SessionModelExecution(t.Context(), tenant, stale.ID) if err != nil || frozen == nil || *frozen != *before.Provider { t.Fatal("frozen tuple bundle changed", err) } @@ -253,7 +253,7 @@ func TestSessionCreationKeepsItsResolvedDeploymentRevision(t *testing.T) { if recreated := resolve(); recreated.Revision == original || recreated.Revision == current.Revision { t.Fatal("revision reused") } - frozenProvider, err := s.SessionModelExecution(t.Context(), tenant, session.ID) + frozenProvider, err := sessionAdapter(s).SessionModelExecution(t.Context(), tenant, session.ID) if err != nil || frozenProvider == nil || *frozenProvider != *input.ModelProvider { t.Fatal("replacement changed the Session bundle", err) } diff --git a/services/core/internal/store/session_creation_stream_test.go b/services/core/internal/store/session_creation_stream_test.go index 6c63da92e..ce7b03c1d 100644 --- a/services/core/internal/store/session_creation_stream_test.go +++ b/services/core/internal/store/session_creation_stream_test.go @@ -137,7 +137,7 @@ func TestCreationStreamStartsBeforeOwnInputsAndRetriesAtUpsertCursor(t *testing. t.Fatal("later Turn events are out of order", i, future[i].Event.Type) } } - turns, err := s.ListTurns(ctx, tenant, id, "", 100, true) + turns, err := sessionAdapter(s).ListTurns(ctx, tenant, id, "", 100, true) if err != nil || len(turns.Turns) != 2 { t.Fatal(turns, err) } diff --git a/services/core/internal/store/session_deletion_execution_test.go b/services/core/internal/store/session_deletion_execution_test.go index 83044b453..678580d56 100644 --- a/services/core/internal/store/session_deletion_execution_test.go +++ b/services/core/internal/store/session_deletion_execution_test.go @@ -9,6 +9,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestDeletedSessionWaitingTurnSettlesWithoutStoppingWorker(t *testing.T) { @@ -47,7 +48,7 @@ func TestDeletedSessionWaitingTurnSettlesWithoutStoppingWorker(t *testing.T) { } h.write(input.TurnID, proto.TypeInteractionDecisionAck, proto.InteractionDecisionAckPayload{DeliveryID: request.DeliveryID, Applied: true, Outcome: &proto.DonePayload{Metadata: map[string]any{proto.DoneMetaAgentSessionID: "deleted-native"}}}) waitTurn(t, h, input.TurnID, sessions.TurnCancelled) - bound, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) + bound, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || bound.NativeSessionID != "deleted-native" { t.Fatal(bound, err) } @@ -68,7 +69,7 @@ func TestDeletedSessionRestartStillReconcilesHiddenClaim(t *testing.T) { h := newDispatchHarness(t) input := h.message("interrupted", "Run") ctx := t.Context() - if _, err := h.s.TransitionTurn(ctx, h.tenant, h.session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err := store.TransitionTurn(ctx, h.s, h.tenant, h.session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } if err := h.s.DeleteSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotIdle) { @@ -83,7 +84,7 @@ func TestDeletedSessionRestartStillReconcilesHiddenClaim(t *testing.T) { if err := worker.Run(stopped); !errors.Is(err, context.Canceled) { t.Fatal(err) } - turn, err := h.s.GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnFailed { t.Fatal(turn, err) } diff --git a/services/core/internal/store/session_deletion_lifecycle_public_test.go b/services/core/internal/store/session_deletion_lifecycle_public_test.go index 5eaf767e5..04fd25a7f 100644 --- a/services/core/internal/store/session_deletion_lifecycle_public_test.go +++ b/services/core/internal/store/session_deletion_lifecycle_public_test.go @@ -73,9 +73,9 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { case "function": err = leased.Sessions.RecordFunctionCall(ctx, tenant, session.ID, receipt.TurnID, sessions.FunctionCall{CallID: "pending", ExecutorCallID: "native-pending", Name: "lookup", Arguments: json.RawMessage(`{}`)}) case sessions.TurnCompleted, sessions.TurnFailed: - _, err = s.CompleteExecution(ctx, tenant, session.ID, receipt.TurnID, status, nil, "", receipt.Sequence) + _, err = leased.Sessions.CompleteExecution(ctx, tenant, session.ID, receipt.TurnID, status, nil, "", receipt.Sequence) default: - _, err = s.TransitionTurn(ctx, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: from, Status: status}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: from, Status: status}) from = status } if err != nil { diff --git a/services/core/internal/store/session_deletion_test.go b/services/core/internal/store/session_deletion_test.go index 84b0d53f5..f1bd0de96 100644 --- a/services/core/internal/store/session_deletion_test.go +++ b/services/core/internal/store/session_deletion_test.go @@ -67,13 +67,13 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { t.Fatal(err) } if status != sessions.TurnQueued { - _, err = s.TransitionTurn(ctx, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = transitionTurn(ctx, s, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } } if status == sessions.TurnCompleted || status == sessions.TurnFailed { - if _, err = s.CompleteExecution(ctx, tenant, session.ID, receipt.TurnID, status, nil, "", receipt.Sequence); err != nil { + if _, err = completeExecution(ctx, t, s, tenant, session.ID, receipt.TurnID, status, nil, "", receipt.Sequence); err != nil { t.Fatal(err) } } @@ -89,7 +89,7 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { if err := s.DeleteSession(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("active Session deleted", err) } - turn, err := s.GetTurn(ctx, tenant, session.ID, receipt.TurnID) + turn, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, receipt.TurnID) if err != nil || turn.Status != status || !turn.CancelRequestedAt.IsZero() { t.Fatal("rejected deletion changed the Turn", turn, err) } @@ -108,7 +108,7 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { if err := s.DeleteSession(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("cancelling Session deleted", err) } - if _, err := s.CompleteExecution(ctx, tenant, session.ID, receipt.TurnID, sessions.TurnCancelled, nil, "", receipt.Sequence); err != nil { + if _, err := completeExecution(ctx, t, s, tenant, session.ID, receipt.TurnID, sessions.TurnCancelled, nil, "", receipt.Sequence); err != nil { t.Fatal(err) } } @@ -148,7 +148,7 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { if _, err := sessionAdapter(fresh).ListItems(ctx, tenant, session.ID, "", 20, true); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - turn, err := fresh.GetTurn(ctx, tenant, session.ID, receipt.TurnID) + turn, err := sessionAdapter(fresh).GetTurn(ctx, tenant, session.ID, receipt.TurnID) if err != nil { t.Fatal(err) } @@ -159,7 +159,7 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { if turn.Status != want { t.Fatal(turn) } - if _, err := fresh.TransitionTurn(ctx, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := transitionTurn(ctx, fresh, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal(err) } inputs, err := fresh.ListTurnInputs(ctx, tenant, session.ID, receipt.TurnID, 0, 20) diff --git a/services/core/internal/store/session_diagnostics.go b/services/core/internal/store/session_diagnostics.go index 316c65a21..0c9ee313a 100644 --- a/services/core/internal/store/session_diagnostics.go +++ b/services/core/internal/store/session_diagnostics.go @@ -42,7 +42,7 @@ func (s *Store) GetSessionDiagnosticsSnapshot(ctx context.Context, tenantID, ses // GetTurnDiagnosticsSnapshot reads only root Turns and their root Items. Its // bounded query and public projection share the same committed snapshot. func (s *Store) GetTurnDiagnosticsSnapshot(ctx context.Context, tenantID, sessionID, turnID string) (sessions.TurnDiagnosticsSnapshot, error) { - params, err := publicTurnLookup(tenantID, sessionID, turnID) + params, err := sessionpg.PublicTurnLookup(tenantID, sessionID, turnID) if err != nil { return sessions.TurnDiagnosticsSnapshot{}, err } diff --git a/services/core/internal/store/session_diagnostics_test.go b/services/core/internal/store/session_diagnostics_test.go index 5de8dba1f..9757c2dfe 100644 --- a/services/core/internal/store/session_diagnostics_test.go +++ b/services/core/internal/store/session_diagnostics_test.go @@ -76,7 +76,7 @@ func TestDiagnosticForceSettlementIgnoresNativeClock(t *testing.T) { } source := runtimeDatabaseTime(t, s).Add(skew).UnixMilli() before := runtimeDatabaseTime(t, s) - completed, err := w.CompleteExecution(t.Context(), owner.TenantID, owner.SessionID, turn, sessions.TurnCompleted, json.RawMessage(fmt.Sprintf(`{"done":{"source_completed_at_ms":%d}}`, source)), "", 0) + completed, err := completeExecution(t.Context(), t, w, owner.TenantID, owner.SessionID, turn, sessions.TurnCompleted, json.RawMessage(fmt.Sprintf(`{"done":{"source_completed_at_ms":%d}}`, source)), "", 0) after := runtimeDatabaseTime(t, s) if err != nil || completed.CompletedAt.UnixMilli() != source { t.Fatal("public native completion changed", completed, err) @@ -112,7 +112,7 @@ func TestDiagnosticSettlementWaitsForSessionLock(t *testing.T) { } done := make(chan error, 1) go func() { - _, err := s.TransitionTurn(t.Context(), tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed}) + _, err := transitionTurn(t.Context(), s, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed}) done <- err }() // Wait for the actual competing transaction to block, not a scheduler delay. diff --git a/services/core/internal/store/session_events_test.go b/services/core/internal/store/session_events_test.go index e65fbf79b..48fa34581 100644 --- a/services/core/internal/store/session_events_test.go +++ b/services/core/internal/store/session_events_test.go @@ -72,7 +72,7 @@ func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { if before != after { t.Fatal("input retry published duplicate events") } - if _, err = s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err = store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } batch := []sessions.ExecutionEvent{ @@ -97,7 +97,7 @@ func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { if before != after { t.Fatal("failed transaction published events") } - if _, err = s.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{"private":"must not escape"}`), "", input.Sequence); err != nil { + if _, err = journal.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{"private":"must not escape"}`), "", input.Sequence); err != nil { t.Fatal(err) } var all []sessions.SessionChange diff --git a/services/core/internal/store/session_metadata_test.go b/services/core/internal/store/session_metadata_test.go index 771ef210d..4e775bda4 100644 --- a/services/core/internal/store/session_metadata_test.go +++ b/services/core/internal/store/session_metadata_test.go @@ -118,10 +118,10 @@ func TestSessionMetadataPreservesTerminalActivity(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.TransitionTurn(ctx, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err := transitionTurn(ctx, s, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } - if _, err := s.TransitionTurn(ctx, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: status}); err != nil { + if _, err := transitionTurn(ctx, s, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: status}); err != nil { t.Fatal(err) } before, err := s.GetSession(ctx, tenant, session.ID) diff --git a/services/core/internal/store/session_model_execution.go b/services/core/internal/store/session_model_execution.go index ad6565c49..c3f9d469b 100644 --- a/services/core/internal/store/session_model_execution.go +++ b/services/core/internal/store/session_model_execution.go @@ -1,11 +1,8 @@ package store import ( - "bytes" "context" "encoding/json" - "errors" - "io" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" @@ -28,32 +25,3 @@ func (s *Store) saveSessionModelExecution(ctx context.Context, q *sqlc.Queries, } return q.SaveSessionModelExecution(ctx, sqlc.SaveSessionModelExecutionParams{SessionID: session, EncryptedConfig: encrypted}) } - -func (s *Store) SessionModelExecution(ctx context.Context, tenant, session string) (*v1.ModelProviderInput, error) { - tenantID, err := parseID(tenant) - if err != nil { - return nil, err - } - sessionID, err := parseID(session) - if err != nil { - return nil, err - } - ciphertext, err := s.queries.GetSessionModelExecution(ctx, sqlc.GetSessionModelExecutionParams{TenantID: tenantID, SessionID: sessionID}) - if err != nil { - return nil, errors.New("session model execution configuration is unavailable") - } - raw, err := s.credentialCipher.OpenModelExecution(ciphertext, tenant, session) - if err != nil { - return nil, errors.New("session model execution decryption is unavailable") - } - var provider v1.ModelProviderInput - decoder := json.NewDecoder(bytes.NewReader(raw)) - decoder.DisallowUnknownFields() - if decoder.Decode(&provider) != nil || decoder.Decode(new(any)) != io.EOF { - return nil, errors.New("invalid stored model execution configuration") - } - if err := provider.Validate(); err != nil { - return nil, err - } - return &provider, nil -} diff --git a/services/core/internal/store/session_model_execution_test.go b/services/core/internal/store/session_model_execution_test.go index e5eb6ec92..303817845 100644 --- a/services/core/internal/store/session_model_execution_test.go +++ b/services/core/internal/store/session_model_execution_test.go @@ -42,18 +42,18 @@ func TestSessionModelExecutionEncryptedAndBound(t *testing.T) { t.Fatal("changed credentials accepted", err) } restarted := withPlacement(t, NewWithCredentialCipher(pool, cipher)) - provider, err := restarted.SessionModelExecution(ctx, tenant, session.ID) + provider, err := sessionAdapter(restarted).SessionModelExecution(ctx, tenant, session.ID) if err != nil || provider.APIKey != "private-model-canary" { t.Fatal("restart lost model credential", err) } - if _, err := restarted.SessionModelExecution(ctx, uuid.NewString(), session.ID); err == nil { - t.Fatal("foreign tenant read credential") + if _, err := sessionAdapter(restarted).SessionModelExecution(ctx, uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal("foreign tenant read credential", err) } if _, err := cipher.OpenModelExecution(ciphertext, tenant, uuid.NewString()); err == nil { t.Fatal("ciphertext not Session-bound") } - if _, err := New(pool).SessionModelExecution(ctx, tenant, session.ID); err == nil { - t.Fatal("missing cipher succeeded") + if _, err := sessionAdapter(New(pool)).SessionModelExecution(ctx, tenant, session.ID); !errors.Is(err, credentialcrypto.ErrUnavailable) { + t.Fatal("missing cipher succeeded", err) } input.IdempotencyKey = uuid.NewString() if _, err := withPlacement(t, New(pool)).CreateSession(ctx, tenant, input); !errors.Is(err, credentialcrypto.ErrUnavailable) { diff --git a/services/core/internal/store/session_model_options_test.go b/services/core/internal/store/session_model_options_test.go index 19e76c80c..079b3b06e 100644 --- a/services/core/internal/store/session_model_options_test.go +++ b/services/core/internal/store/session_model_options_test.go @@ -25,7 +25,7 @@ func TestSessionModelExecutionStoresOnlyProviderBundle(t *testing.T) { if err != nil { t.Fatal(err) } - flat, err := st.SessionModelExecution(ctx, tenant, session.ID) + flat, err := sessionAdapter(st).SessionModelExecution(ctx, tenant, session.ID) if err != nil || flat == nil || *flat != *provider { t.Fatal("new Session did not store a flat provider bundle", err) } @@ -47,7 +47,7 @@ func TestSessionModelExecutionStoresOnlyProviderBundle(t *testing.T) { if _, err := pool.Exec(ctx, "UPDATE session_model_execution SET encrypted_config=$2 WHERE session_id=$1", session.ID, ciphertext); err != nil { t.Fatal(err) } - if _, err := NewWithCredentialCipher(pool, cipher).SessionModelExecution(ctx, tenant, session.ID); err == nil { + if _, err := sessionAdapter(NewWithCredentialCipher(pool, cipher)).SessionModelExecution(ctx, tenant, session.ID); err == nil { t.Fatal("retired native options accepted in provider bundle") } } diff --git a/services/core/internal/store/structured_output_dispatch_test.go b/services/core/internal/store/structured_output_dispatch_test.go index 86bc24e64..94b6e34f9 100644 --- a/services/core/internal/store/structured_output_dispatch_test.go +++ b/services/core/internal/store/structured_output_dispatch_test.go @@ -54,7 +54,7 @@ func TestStructuredOutputDispatchRechecksOperationQualification(t *testing.T) { if result := <-h.run(ctx, input.TurnID); result.err == nil { t.Fatal("unqualified structured output dispatched") } - turn, err := h.s.GetTurn(t.Context(), h.tenant, h.session.ID, input.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, input.TurnID) if err != nil || turn.Status != sessions.TurnQueued { t.Fatal("unqualified work was claimed", turn.Status, err) } diff --git a/services/core/internal/store/structured_output_native_test.go b/services/core/internal/store/structured_output_native_test.go index 58d61fd4f..2d51e2995 100644 --- a/services/core/internal/store/structured_output_native_test.go +++ b/services/core/internal/store/structured_output_native_test.go @@ -74,7 +74,7 @@ func TestNativeStructuredOutputPublicExecution(t *testing.T) { if err != nil || !call.Applied { t.Fatal("function application receipt missing", err) } - turn, err := h.s.GetTurn(ctx, h.tenant, proof.Session, proof.Turn) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, proof.Session, proof.Turn) if err != nil { t.Fatal(err) } @@ -82,14 +82,14 @@ func TestNativeStructuredOutputPublicExecution(t *testing.T) { if json.Unmarshal(turn.Outcome, &outcome) != nil || outcome.Done.Usage.Raw["claude_sdk_result"] == nil || outcome.AppliedThrough < 1 { t.Fatal("native usage or input receipt missing") } - before, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + before, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID == "" { t.Fatal("native binding missing", err) } stop() stop = startNativeEngineDaemon(t, h, home, binary, "claude_sdk") run("resume") - after, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + after, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID != after.NativeSessionID { t.Fatal("native history changed", err) } diff --git a/services/core/internal/store/subagent_identities_test.go b/services/core/internal/store/subagent_identities_test.go index 03857d54f..f5dc4a1a5 100644 --- a/services/core/internal/store/subagent_identities_test.go +++ b/services/core/internal/store/subagent_identities_test.go @@ -103,7 +103,7 @@ func TestSubagentIdentityIsAtomicScopedAndImmutable(t *testing.T) { if err = journal.AppendTurnEvents(ctx, tenant, foreign.ID, foreignInput.TurnID, 1, []sessions.ExecutionEvent{subagentIdentityEvent("other-child", "root", 101)}); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal("known root binding ignored", err) } - if _, err = w.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCompleted, json.RawMessage(`{}`), "root", input.Sequence); err != nil { + if _, err = completeExecution(ctx, t, w, tenant, session.ID, input.TurnID, sessions.TurnCompleted, json.RawMessage(`{}`), "root", input.Sequence); err != nil { t.Fatal(err) } awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, w.pool) diff --git a/services/core/internal/store/subagent_native_outputs_test.go b/services/core/internal/store/subagent_native_outputs_test.go index caaf9946a..f6154fb3b 100644 --- a/services/core/internal/store/subagent_native_outputs_test.go +++ b/services/core/internal/store/subagent_native_outputs_test.go @@ -136,13 +136,13 @@ func TestSubagentRootCompletionRetainsNativeSourceTime(t *testing.T) { tenant, session := newTurnSession(t, s) input := submitMessage(t, s, tenant, session.ID, "start") transition(t, s, tenant, session.ID, input.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - current, err := s.GetTurn(t.Context(), tenant, session.ID, input.TurnID) + current, err := sessionAdapter(s).GetTurn(t.Context(), tenant, session.ID, input.TurnID) if err != nil { t.Fatal(err) } source := current.CreatedAt.Unix() * 1000 outcome := json.RawMessage(fmt.Sprintf(`{"done":{"source_completed_at_ms":%d}}`, source)) - completed, err := s.CompleteExecution(t.Context(), tenant, session.ID, input.TurnID, status, outcome, "", input.Sequence) + completed, err := completeExecution(t.Context(), t, s, tenant, session.ID, input.TurnID, status, outcome, "", input.Sequence) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/subagent_resources_test.go b/services/core/internal/store/subagent_resources_test.go index 49eecf66f..87135d8a1 100644 --- a/services/core/internal/store/subagent_resources_test.go +++ b/services/core/internal/store/subagent_resources_test.go @@ -90,13 +90,13 @@ func TestSubagentResourcesNativeOwnershipLifecycleAndRecovery(t *testing.T) { t.Fatal(same, err) } // Session Turn reads carry root work only: a child Turn ID is missing there. - if _, err = s.GetTurn(ctx, tenant, session.ID, tid); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).GetTurn(ctx, tenant, session.ID, tid); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("child Turn in Session Turn retrieval", err) } - if _, err = s.ListTurns(ctx, tenant, session.ID, tid, 100, true); !errors.Is(err, sessions.ErrNotFound) { + if _, err = sessionAdapter(s).ListTurns(ctx, tenant, session.ID, tid, 100, true); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("child Turn as a Session Turn cursor", err) } - allTurns, err := s.ListTurns(ctx, tenant, session.ID, "", 100, true) + allTurns, err := sessionAdapter(s).ListTurns(ctx, tenant, session.ID, "", 100, true) if err != nil || len(allTurns.Turns) != 1 || allTurns.Turns[0].ID != root.TurnID { t.Fatal(allTurns, err) } @@ -113,7 +113,7 @@ func TestSubagentResourcesNativeOwnershipLifecycleAndRecovery(t *testing.T) { if _, err = sessionAdapter(s).ListSubagentTurns(ctx, uuid.NewString(), session.ID, child.ID, "", 20, true); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign tenant child Turns", err) } - if allTurns, err = s.ListTurns(ctx, tenant, session.ID, "", 100, true); err != nil || len(allTurns.Turns) != 1 { + if allTurns, err = sessionAdapter(s).ListTurns(ctx, tenant, session.ID, "", 100, true); err != nil || len(allTurns.Turns) != 1 { t.Fatal(allTurns, err) } own, err := sessionAdapter(s).ListSubagentTurnItems(ctx, tenant, session.ID, child.ID, tid, "", 20, true) diff --git a/services/core/internal/store/subagent_visibility_public_test.go b/services/core/internal/store/subagent_visibility_public_test.go index 80bc82abd..4888d5895 100644 --- a/services/core/internal/store/subagent_visibility_public_test.go +++ b/services/core/internal/store/subagent_visibility_public_test.go @@ -11,6 +11,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -115,7 +116,7 @@ func TestSubagentVisibilityPublic(t *testing.T) { } observed := collectEvents(t, live) - page, err := s.ListTurns(ctx, tenant, session, "", 10, true) + page, err := store.SessionAdapter(s).ListTurns(ctx, tenant, session, "", 10, true) if err != nil || len(page.Turns) != 1 { t.Fatal(page, err) } @@ -127,7 +128,7 @@ func TestSubagentVisibilityPublic(t *testing.T) { if err = leased.Sessions.BindSessionDevice(ctx, tenant, session, host.ID); err != nil { t.Fatal(err) } - if _, err = writer.TransitionTurn(ctx, tenant, session, root, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if _, err = store.TransitionTurn(ctx, writer, tenant, session, root, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } identity := func(child, parent string, created int64) sessions.ExecutionEvent { @@ -151,7 +152,7 @@ func TestSubagentVisibilityPublic(t *testing.T) { if err = leased.Sessions.AppendTurnEvents(ctx, tenant, session, root, 1, facts); err != nil { t.Fatal(err) } - if _, err = writer.TransitionTurn(ctx, tenant, session, root, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnCompleted}); err != nil { + if _, err = store.TransitionTurn(ctx, writer, tenant, session, root, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnCompleted}); err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/token_usage_integration_test.go b/services/core/internal/store/token_usage_integration_test.go index 404c8a8a9..2dd7268ff 100644 --- a/services/core/internal/store/token_usage_integration_test.go +++ b/services/core/internal/store/token_usage_integration_test.go @@ -38,7 +38,7 @@ func TestTokenUsageDurableSnapshotsAndSessionTotals(t *testing.T) { if err != nil { t.Fatal(err) } - _, err = s.TransitionTurn(ctx, tenant, session.ID, admission.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, admission.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } @@ -52,29 +52,29 @@ func TestTokenUsageDurableSnapshotsAndSessionTotals(t *testing.T) { if err = journal.AppendTurnEvents(ctx, tenant, session.ID, admission.TurnID, 2, []sessions.ExecutionEvent{{Kind: "usage", Payload: usage(20)}}); err != nil { t.Fatal(err) } - measured, err := s.GetTurn(ctx, tenant, session.ID, admission.TurnID) + measured, err := store.SessionAdapter(s).GetTurn(ctx, tenant, session.ID, admission.TurnID) if err != nil { t.Fatal(err) } check(measured.Usage, 20) - if _, err = s.CompleteExecution(ctx, tenant, session.ID, admission.TurnID, sessions.TurnCompleted, json.RawMessage(`{"done":{"usage":`+string(usage(99))+`}}`), "missing-binding", admission.Sequence); !errors.Is(err, sessions.ErrNotFound) { + if _, err = journal.CompleteExecution(ctx, tenant, session.ID, admission.TurnID, sessions.TurnCompleted, json.RawMessage(`{"done":{"usage":`+string(usage(99))+`}}`), "missing-binding", admission.Sequence); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - rolledBack, err := s.GetTurn(ctx, tenant, session.ID, admission.TurnID) + rolledBack, err := store.SessionAdapter(s).GetTurn(ctx, tenant, session.ID, admission.TurnID) if err != nil || rolledBack.Status != sessions.TurnInProgress { t.Fatalf("rollback: %+v %v", rolledBack, err) } check(rolledBack.Usage, 20) // Completion without usage retains the last persisted measurement. - completed, err := s.CompleteExecution(ctx, tenant, session.ID, admission.TurnID, status, json.RawMessage(`{"done":{"content":"partial"}}`), "", admission.Sequence) + completed, err := journal.CompleteExecution(ctx, tenant, session.ID, admission.TurnID, status, json.RawMessage(`{"done":{"content":"partial"}}`), "", admission.Sequence) if err != nil { t.Fatal(err) } check(completed.Usage, 20) - if _, err = s.CompleteExecution(ctx, tenant, session.ID, admission.TurnID, status, json.RawMessage(`{"done":{"usage":`+string(usage(99))+`}}`), "", admission.Sequence); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err = journal.CompleteExecution(ctx, tenant, session.ID, admission.TurnID, status, json.RawMessage(`{"done":{"usage":`+string(usage(99))+`}}`), "", admission.Sequence); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal(err) } - if _, err = s.GetTurn(ctx, uuid.NewString(), session.ID, admission.TurnID); !errors.Is(err, sessions.ErrNotFound) { + if _, err = store.SessionAdapter(s).GetTurn(ctx, uuid.NewString(), session.ID, admission.TurnID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } } @@ -118,7 +118,7 @@ func TestCancellationReceiptUsageSurvivesRecovery(t *testing.T) { if err != nil { t.Fatal(err) } - _, err = s.TransitionTurn(ctx, tenant, session.ID, admission.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, admission.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } @@ -127,7 +127,7 @@ func TestCancellationReceiptUsageSurvivesRecovery(t *testing.T) { t.Fatal(err) } // Startup recovery has no in-memory cancellation outcome. - recovered, err := s.TransitionTurn(ctx, tenant, session.ID, admission.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"error_code":"execution_interrupted"}`)}) + recovered, err := store.TransitionTurn(ctx, s, tenant, session.ID, admission.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"error_code":"execution_interrupted"}`)}) if err != nil { t.Fatal(err) } @@ -196,13 +196,13 @@ func TestSessionUsageRequiresEveryRootTurnEndedAndMeasured(t *testing.T) { } move := func(turn, from, to string) { t.Helper() - if _, err := s.TransitionTurn(ctx, tenant, session.ID, turn, sessions.TurnTransition{ExpectedStatus: from, Status: to}); err != nil { + if _, err := store.TransitionTurn(ctx, s, tenant, session.ID, turn, sessions.TurnTransition{ExpectedStatus: from, Status: to}); err != nil { t.Fatal(err) } } finish := func(admission sessions.InputReceipt, status string) { t.Helper() - if _, err := s.CompleteExecution(ctx, tenant, session.ID, admission.TurnID, status, json.RawMessage(`{"done":{}}`), "", admission.Sequence); err != nil { + if _, err := journal.CompleteExecution(ctx, tenant, session.ID, admission.TurnID, status, json.RawMessage(`{"done":{}}`), "", admission.Sequence); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/tool_policy_native_test.go b/services/core/internal/store/tool_policy_native_test.go index aea5f0195..821c075cf 100644 --- a/services/core/internal/store/tool_policy_native_test.go +++ b/services/core/internal/store/tool_policy_native_test.go @@ -12,6 +12,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -90,7 +91,7 @@ func TestNativeToolPolicyPublicExecution(t *testing.T) { if err != nil || session.Engine != kind { t.Fatal("selected engine was not persisted", err) } - turn, err := h.s.GetTurn(ctx, h.tenant, item.ID, item.FirstTurn) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, item.ID, item.FirstTurn) if err != nil { t.Fatal(err) } @@ -99,7 +100,7 @@ func TestNativeToolPolicyPublicExecution(t *testing.T) { if err != nil || len(inputs) != 1 || json.Unmarshal(turn.Outcome, &outcome) != nil || outcome.AppliedThrough != inputs[0].Sequence { t.Fatal("native text input receipt missing", err) } - binding, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, item.ID) + binding, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, item.ID) if err != nil || binding.NativeSessionID == "" { t.Fatal("native binding missing", err) } @@ -109,7 +110,7 @@ func TestNativeToolPolicyPublicExecution(t *testing.T) { stop = startNativeEngineDaemon(t, h, home, binary, kind) run("resume") for id, before := range nativeIDs { - after, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, id) + after, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, id) if err != nil || after.NativeSessionID != before { t.Fatal("cold continuation changed native history", err) } diff --git a/services/core/internal/store/tool_search_native_test.go b/services/core/internal/store/tool_search_native_test.go index e2123d634..626443f9b 100644 --- a/services/core/internal/store/tool_search_native_test.go +++ b/services/core/internal/store/tool_search_native_test.go @@ -74,7 +74,7 @@ func TestNativeToolSearchPublicExecution(t *testing.T) { if err != nil || !call.Applied { t.Fatal("function application receipt missing", err) } - turn, err := h.s.GetTurn(ctx, h.tenant, proof.Session, proof.Turn) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, proof.Session, proof.Turn) if err != nil { t.Fatal(err) } @@ -82,14 +82,14 @@ func TestNativeToolSearchPublicExecution(t *testing.T) { if json.Unmarshal(turn.Outcome, &outcome) != nil || outcome.Done.Usage.Raw["claude_sdk_result"] == nil || outcome.AppliedThrough < 1 { t.Fatal("native usage or input receipt missing") } - before, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + before, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID == "" { t.Fatal("native binding missing", err) } stop() stop = startNativeEngineDaemon(t, h, home, binary, "claude_sdk") run("resume") - after, err := h.s.GetSessionExecutionBinding(ctx, h.tenant, proof.Session) + after, err := store.SessionAdapter(h.s).GetSessionExecutionBinding(ctx, h.tenant, proof.Session) if err != nil || before.NativeSessionID != after.NativeSessionID { t.Fatal("native history changed", err) } diff --git a/services/core/internal/store/turn_completion.go b/services/core/internal/store/turn_completion.go deleted file mode 100644 index bd7f812ca..000000000 --- a/services/core/internal/store/turn_completion.go +++ /dev/null @@ -1,105 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" - "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" -) - -// CompleteExecution commits the outcome and native continuity under the admission lock. -func (s *Store) CompleteExecution(ctx context.Context, tenantID, sessionID, turnID, status string, outcome json.RawMessage, nativeID string, appliedThrough int64) (sessions.Turn, error) { - p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) - if err != nil { - return sessions.Turn{}, err - } - if !sessions.TerminalStatus(status) || len(outcome) > 512*1024 || len(nativeID) > 512 || appliedThrough < 0 { - return sessions.Turn{}, sessions.ErrInvalidInput - } - outcome, err = jsonobject.Normalize(outcome) - if err != nil { - return sessions.Turn{}, fmt.Errorf("%w: %w", sessions.ErrInvalidInput, err) - } - var row sqlc.Turn - err = s.withSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - current, err := q.GetTurn(ctx, p) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } else if err != nil { - return err - } - if current.Status != sessions.TurnInProgress && (current.Status != sessions.TurnWaiting || status == sessions.TurnCompleted) { - return sessions.ErrTurnConflict - } - if status == sessions.TurnCompleted { - pending, err := q.HasUnappliedMessages(ctx, sqlc.HasUnappliedMessagesParams{SessionID: session, TurnID: p.ID, Sequence: appliedThrough}) - if err != nil { - return err - } - if pending { - return sessions.ErrUnappliedInputs - } - } - sourceCompleted := pgtype.Timestamptz{} - if status == sessions.TurnCompleted { - var snapshot struct { - Done *struct { - SourceCompletedAtMS *int64 `json:"source_completed_at_ms"` - } `json:"done"` - } - if json.Unmarshal(outcome, &snapshot) != nil { - return sessions.ErrInvalidInput - } - if snapshot.Done != nil && snapshot.Done.SourceCompletedAtMS != nil { - ms := *snapshot.Done.SourceCompletedAtMS - if ms <= 0 { - return sessions.ErrInvalidInput - } - // Native and Core timestamps come from independent host clocks. - // Preserve source time; committed activity uses the database clock. - sourceCompleted = pgtype.Timestamptz{Time: time.UnixMilli(ms), Valid: true} - } - } - row, err = q.TransitionTurn(ctx, sqlc.TransitionTurnParams{ID: p.ID, SessionID: session, ExpectedStatus: current.Status, NewStatus: status, Outcome: outcome, SourceCompletedAt: sourceCompleted}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrTurnConflict - } - if err != nil { - return err - } - if err = sessions.AppendTurnEvent(ctx, sessionpg.BindSession(q, p.TenantID, session), uuid.UUID(row.ID.Bytes).String(), row.EventCount, sessions.ExecutionEvent{Kind: "execution_" + status, Payload: outcome}); err != nil { - return err - } - if nativeID != "" { - n, err := q.RememberNativeSession(ctx, sqlc.RememberNativeSessionParams{SessionID: session, NativeSessionID: nativeID}) - if err != nil { - return err - } - if n != 1 { - return sessions.ErrNotFound - } - } - row, err = q.GetTurn(ctx, p) - if err != nil { - return err - } - ending, err := sessionpg.LoadEnding(ctx, q, session, row.ID) - if err != nil { - return err - } - return sessionpg.ApplyTurnEnd(ctx, q, session, row.ID, sessions.EndTurn(sessionpg.TurnFromRow(row), ending)) - }) - if err != nil { - return sessions.Turn{}, err - } - return sessionpg.TurnFromRow(row), nil -} diff --git a/services/core/internal/store/turn_events_test.go b/services/core/internal/store/turn_events_test.go index cbdd71512..bd4c1c9ec 100644 --- a/services/core/internal/store/turn_events_test.go +++ b/services/core/internal/store/turn_events_test.go @@ -25,7 +25,7 @@ func TestTurnEventBatchesAreOrderedIsolatedAndDurable(t *testing.T) { if err != nil { t.Fatal(err) } - _, err = s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err = store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } @@ -66,14 +66,14 @@ func TestTurnEventBatchesAreOrderedIsolatedAndDurable(t *testing.T) { t.Fatal(err) } // A failed native binding write must roll back both the terminal event and status. - if _, err = s.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCompleted, json.RawMessage(`{}`), "missing-binding", input.Sequence); !errors.Is(err, sessions.ErrNotFound) { + if _, err = journal.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCompleted, json.RawMessage(`{}`), "missing-binding", input.Sequence); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } events, _ := s.ListTurnEvents(ctx, tenant, session.ID, input.TurnID, 0, 100) if len(events) != 2 { t.Fatal("terminal event survived rollback") } - if _, err = s.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{"done":{"content":""}}`), "", input.Sequence); err != nil { + if _, err = journal.CompleteExecution(ctx, tenant, session.ID, input.TurnID, sessions.TurnCancelled, json.RawMessage(`{"done":{"content":""}}`), "", input.Sequence); err != nil { t.Fatal(err) } if err = journal.AppendTurnEvents(ctx, tenant, session.ID, input.TurnID, 1, batch); err != nil { @@ -98,7 +98,7 @@ func TestEventLimitStillAllowsTerminalFailure(t *testing.T) { h := newDispatchHarness(t) ctx := context.Background() input := h.message("start", "Test output budget") - _, err := h.s.TransitionTurn(ctx, h.tenant, h.session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) + _, err := store.TransitionTurn(ctx, h.s, h.tenant, h.session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) if err != nil { t.Fatal(err) } @@ -111,7 +111,7 @@ func TestEventLimitStillAllowsTerminalFailure(t *testing.T) { if err = h.owner().Sessions.AppendTurnEvents(ctx, h.tenant, h.session.ID, input.TurnID, 1, events); !errors.Is(err, sessions.ErrEventLimit) { t.Fatal(err) } - if _, err = h.s.CompleteExecution(ctx, h.tenant, h.session.ID, input.TurnID, sessions.TurnFailed, json.RawMessage(`{"error_code":"event_limit"}`), "", input.Sequence); err != nil { + if _, err = h.owner().Sessions.CompleteExecution(ctx, h.tenant, h.session.ID, input.TurnID, sessions.TurnFailed, json.RawMessage(`{"error_code":"event_limit"}`), "", input.Sequence); err != nil { t.Fatal(err) } got, err := h.s.ListTurnEvents(ctx, h.tenant, h.session.ID, input.TurnID, 0, 100) diff --git a/services/core/internal/store/turn_inputs.go b/services/core/internal/store/turn_inputs.go index c9f8285ce..3e896dafd 100644 --- a/services/core/internal/store/turn_inputs.go +++ b/services/core/internal/store/turn_inputs.go @@ -211,8 +211,10 @@ func (s *Store) ListTurnInputs(ctx context.Context, tenantID, sessionID, turnID if after < 0 || limit < 1 || limit > 100 { return nil, fmt.Errorf("%w: nonnegative cursor and page size 1..100 required", sessions.ErrInvalidInput) } - if _, err := s.GetTurn(ctx, tenantID, sessionID, turnID); err != nil { - return nil, err + if _, err := s.queries.GetTurn(ctx, params); errors.Is(err, pgx.ErrNoRows) { + return nil, sessions.ErrNotFound + } else if err != nil { + return nil, fmt.Errorf("get turn: %w", err) } rows, err := s.queries.ListTurnInputs(ctx, sqlc.ListTurnInputsParams{ TenantID: params.TenantID, SessionID: params.SessionID, TurnID: params.ID, Sequence: after, Limit: int32(limit), diff --git a/services/core/internal/store/turn_inputs_test.go b/services/core/internal/store/turn_inputs_test.go index aee9a919b..c5924e554 100644 --- a/services/core/internal/store/turn_inputs_test.go +++ b/services/core/internal/store/turn_inputs_test.go @@ -38,7 +38,7 @@ func submitMessage(t *testing.T, s *Store, tenant, session, key string) sessions func transition(t *testing.T, s *Store, tenant, session, turn, from, to string) sessions.Turn { t.Helper() - got, err := s.TransitionTurn(context.Background(), tenant, session, turn, sessions.TurnTransition{ExpectedStatus: from, Status: to}) + got, err := transitionTurn(context.Background(), s, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: from, Status: to}) if err != nil { t.Fatal(err) } @@ -137,7 +137,7 @@ func TestTurnInputRetriesAndRestart(t *testing.T) { if !retry.Replayed || retry.TurnID != first.TurnID || retry.Sequence != first.Sequence { t.Fatalf("restart retry changed target: %+v", retry) } - got, err := recovered.GetTurn(ctx, tenant, session.ID, first.TurnID) + got, err := sessionAdapter(recovered).GetTurn(ctx, tenant, session.ID, first.TurnID) if err != nil || !reflect.DeepEqual(got, completed) { t.Fatalf("restart turn: %+v, %v", got, err) } @@ -183,13 +183,16 @@ func TestTurnOperationsAreTenantAndSessionScoped(t *testing.T) { return err }, "cancel": func() error { _, err := s.RequestCancel(ctx, scope.tenant, scope.session, "cancel"); return err }, - "read": func() error { _, err := s.GetTurn(ctx, scope.tenant, scope.session, first.TurnID); return err }, + "read": func() error { + _, err := sessionAdapter(s).GetTurn(ctx, scope.tenant, scope.session, first.TurnID) + return err + }, "inputs": func() error { _, err := s.ListTurnInputs(ctx, scope.tenant, scope.session, first.TurnID, 0, 10) return err }, "transition": func() error { - _, err := s.TransitionTurn(ctx, scope.tenant, scope.session, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed}) + _, err := transitionTurn(ctx, s, scope.tenant, scope.session, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed}) return err }, } { @@ -203,10 +206,10 @@ func TestTurnOperationsAreTenantAndSessionScoped(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.GetTurn(ctx, tenant, second.ID, first.TurnID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetTurn(ctx, tenant, second.ID, first.TurnID); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("cross-session turn read: %v", err) } - if _, err := s.TransitionTurn(ctx, tenant, second.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed}); !errors.Is(err, sessions.ErrNotFound) { + if _, err := transitionTurn(ctx, s, tenant, second.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed}); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("cross-session turn write: %v", err) } otherInput := submitMessage(t, s, otherTenant, other.ID, "input") diff --git a/services/core/internal/store/turn_journal_reads_test.go b/services/core/internal/store/turn_journal_reads_test.go index 5ac948041..2e07f9b3d 100644 --- a/services/core/internal/store/turn_journal_reads_test.go +++ b/services/core/internal/store/turn_journal_reads_test.go @@ -23,7 +23,7 @@ func (s *Store) ListTurnEvents(ctx context.Context, tenantID, sessionID, turnID if after < 0 || limit < 1 || limit > 100 { return nil, sessions.ErrInvalidInput } - if _, err = s.GetTurn(ctx, tenantID, sessionID, turnID); err != nil { + if _, err = sessionAdapter(s).GetTurn(ctx, tenantID, sessionID, turnID); err != nil { return nil, err } rows, err := s.queries.ListTurnEvents(ctx, sqlc.ListTurnEventsParams{TenantID: p.TenantID, SessionID: p.SessionID, TurnID: p.ID, Ordinal: after, Limit: int32(limit)}) diff --git a/services/core/internal/store/turn_reads.go b/services/core/internal/store/turn_reads.go deleted file mode 100644 index 85b5c1793..000000000 --- a/services/core/internal/store/turn_reads.go +++ /dev/null @@ -1,49 +0,0 @@ -package store - -import ( - "context" - "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" - "github.com/jackc/pgx/v5/pgtype" -) - -// ListTurns pages a Session's root Turns. Subagent Turns are not Session Turns; -// sessionpg.Store.ListSubagentTurns reads them. -func (s *Store) ListTurns(ctx context.Context, tenantID, sessionID, cursor string, limit int, ascending bool) (sessions.TurnPage, error) { - if limit < 1 || limit > 100 { - return sessions.TurnPage{}, fmt.Errorf("%w: page size must be 1..100", sessions.ErrInvalidInput) - } - if _, err := s.GetSession(ctx, tenantID, sessionID); err != nil { - return sessions.TurnPage{}, err - } - tenant, _ := parseID(tenantID) - session, _ := parseID(sessionID) - params := sqlc.ListRootTurnsParams{TenantID: tenant, SessionID: session, PageLimit: int32(limit + 1), AfterID: pgtype.UUID{Valid: true}, Ascending: ascending} - if cursor != "" { - // A child Turn is not a Session Turn, so its ID is a missing cursor here. - after, err := s.GetTurn(ctx, tenantID, sessionID, pgunit.LookupCursor(cursor)) - if err != nil { - return sessions.TurnPage{}, err - } - params.AfterCreated = pgtype.Timestamptz{Time: after.CreatedAt, Valid: true} - params.AfterID, _ = parseID(after.ID) - } - rows, err := s.queries.ListRootTurns(ctx, params) - if err != nil { - return sessions.TurnPage{}, fmt.Errorf("list turns: %w", err) - } - page := sessions.TurnPage{Turns: make([]sessions.Turn, 0, min(limit, len(rows)))} - if len(rows) > limit { - page.NextCursor = uuid.UUID(rows[limit-1].ID.Bytes).String() - rows = rows[:limit] - } - for _, row := range rows { - page.Turns = append(page.Turns, sessionpg.TurnFromRow(row)) - } - return page, nil -} diff --git a/services/core/internal/store/turn_reads_test.go b/services/core/internal/store/turn_reads_test.go deleted file mode 100644 index 7d0ced7aa..000000000 --- a/services/core/internal/store/turn_reads_test.go +++ /dev/null @@ -1,71 +0,0 @@ -package store - -import ( - "context" - "errors" - "sort" - "testing" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" -) - -func TestTurnPaginationRetainsScopeAndOrder(t *testing.T) { - s, pool := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - ids := make([]string, 0, 4) - for i := 0; i < 4; i++ { - receipt := submitMessage(t, s, tenant, session.ID, uuid.NewString()) - transition(t, s, tenant, session.ID, receipt.TurnID, sessions.TurnQueued, sessions.TurnCancelled) - ids = append(ids, receipt.TurnID) - } - // Equal creation times exercise the ID tie-breaker across page boundaries. - if _, err := pool.Exec(ctx, "UPDATE turns SET created_at=$1 WHERE session_id=$2", time.Unix(1700000000, 0), session.ID); err != nil { - t.Fatal(err) - } - sort.Strings(ids) - s = New(pool) - for _, ascending := range []bool{true, false} { - cursor := "" - for i := 0; i < len(ids); i++ { - page, err := s.ListTurns(ctx, tenant, session.ID, cursor, 1, ascending) - at := i - if !ascending { - at = len(ids) - 1 - i - } - if err != nil || len(page.Turns) != 1 || page.Turns[0].ID != ids[at] { - t.Fatalf("page %d: %+v %v", i, page, err) - } - if (page.NextCursor != "") != (i < len(ids)-1) { - t.Fatalf("incorrect has_more: %+v", page) - } - cursor = page.Turns[0].ID - } - page, err := s.ListTurns(ctx, tenant, session.ID, cursor, 1, ascending) - if err != nil || len(page.Turns) != 0 || page.Turns == nil || page.NextCursor != "" { - t.Fatalf("end: %+v %v", page, err) - } - } - otherTenant, otherSession := newTurnSession(t, s) - sameTenantSession, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString()}) - if err != nil { - t.Fatal(err) - } - for _, scope := range [][2]string{{otherTenant, session.ID}, {tenant, otherSession.ID}, {tenant, uuid.NewString()}, {tenant, sameTenantSession.ID}} { - _, err := s.ListTurns(ctx, scope[0], scope[1], ids[0], 1, true) - if !errors.Is(err, sessions.ErrNotFound) { - t.Fatalf("foreign cursor/session accepted: %v", err) - } - } - empty, err := s.ListTurns(ctx, tenant, sameTenantSession.ID, "", 20, false) - if err != nil || empty.Turns == nil || len(empty.Turns) != 0 { - t.Fatalf("empty session: %+v %v", empty, err) - } - for _, limit := range []int{0, 101} { - if _, err := s.ListTurns(ctx, tenant, session.ID, "", limit, false); !errors.Is(err, sessions.ErrInvalidInput) { - t.Fatalf("limit accepted: %v", err) - } - } -} diff --git a/services/core/internal/store/turns.go b/services/core/internal/store/turns.go deleted file mode 100644 index 1ac432d8d..000000000 --- a/services/core/internal/store/turns.go +++ /dev/null @@ -1,116 +0,0 @@ -package store - -import ( - "context" - "errors" - "fmt" - - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" - "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" -) - -// GetTurn reads a root Turn. A Subagent Turn ID is not found here, exactly like -// a missing one; GetSubagentTurn reads child Turns. -func (s *Store) GetTurn(ctx context.Context, tenantID, sessionID, turnID string) (sessions.Turn, error) { - params, err := publicTurnLookup(tenantID, sessionID, turnID) - if err != nil { - return sessions.Turn{}, err - } - row, err := s.queries.GetTurn(ctx, params) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.Turn{}, sessions.ErrNotFound - } - if err != nil { - return sessions.Turn{}, fmt.Errorf("get turn: %w", err) - } - return sessionpg.TurnFromRow(row), nil -} - -// TransitionTurn is a compare-and-set for execution callbacks. Once terminal, -// a Turn cannot be reopened or have its outcome overwritten, including by retries. -// A dispatcher must claim queued -> in_progress before sending work to a daemon. -func (s *Store) TransitionTurn(ctx context.Context, tenantID, sessionID, turnID string, input sessions.TurnTransition) (sessions.Turn, error) { - params, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) - if err != nil { - return sessions.Turn{}, err - } - if !sessions.ValidTransition(input.ExpectedStatus, input.Status) || len(input.Outcome) > 512*1024 { - return sessions.Turn{}, fmt.Errorf("%w: invalid turn transition or outcome size", sessions.ErrInvalidInput) - } - outcome, err := jsonobject.Normalize(input.Outcome) - if err != nil { - return sessions.Turn{}, fmt.Errorf("%w: %w", sessions.ErrInvalidInput, err) - } - if !sessions.TerminalStatus(input.Status) && string(outcome) != "{}" { - return sessions.Turn{}, fmt.Errorf("%w: outcome requires a terminal status", sessions.ErrInvalidInput) - } - input.Outcome = outcome - var row sqlc.Turn - err = s.withSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, _ pgtype.UUID) error { - var err error - row, err = transitionTurn(ctx, q, params, input) - return err - }) - if err != nil { - return sessions.Turn{}, fmt.Errorf("transition turn: %w", err) - } - return sessionpg.TurnFromRow(row), nil -} - -func transitionTurn(ctx context.Context, q *sqlc.Queries, params sqlc.GetTurnParams, input sessions.TurnTransition) (sqlc.Turn, error) { - if _, err := q.GetTurn(ctx, params); errors.Is(err, pgx.ErrNoRows) { - return sqlc.Turn{}, sessions.ErrNotFound - } else if err != nil { - return sqlc.Turn{}, err - } - if input.ExpectedStatus == sessions.TurnQueued && input.Status == sessions.TurnInProgress { - if err := sessions.CheckComputeAdmission(ctx, sessionpg.BindSession(q, params.TenantID, params.SessionID)); err != nil { - return sqlc.Turn{}, err - } - } - row, err := q.TransitionTurn(ctx, sqlc.TransitionTurnParams{ - ID: params.ID, SessionID: params.SessionID, ExpectedStatus: input.ExpectedStatus, - NewStatus: input.Status, Outcome: input.Outcome, - }) - if errors.Is(err, pgx.ErrNoRows) { - return sqlc.Turn{}, sessions.ErrTurnConflict - } - if err != nil { - return sqlc.Turn{}, err - } - if !sessions.TerminalStatus(row.Status) { - if err := sessionpg.AppendChanges(ctx, q, row.SessionID, sessions.TurnChanges(sessionpg.TurnFromRow(row), false)...); err != nil { - return sqlc.Turn{}, err - } - return row, nil - } - outcome := sessions.Source{Turn: uuid.UUID(row.ID.Bytes).String(), Kind: "execution_" + row.Status, Payload: row.Outcome, CreatedAt: row.CompletedAt.Time} - if err = sessions.ProjectSource(ctx, sessionpg.BindSession(q, params.TenantID, row.SessionID), outcome); err != nil { - return sqlc.Turn{}, err - } - if row, err = q.GetTurn(ctx, params); err != nil { - return sqlc.Turn{}, err - } - ending, err := sessionpg.LoadEnding(ctx, q, row.SessionID, row.ID) - if err != nil { - return sqlc.Turn{}, err - } - if err := sessionpg.ApplyTurnEnd(ctx, q, row.SessionID, row.ID, sessions.EndTurn(sessionpg.TurnFromRow(row), ending)); err != nil { - return sqlc.Turn{}, err - } - return row, nil -} - -// publicTurnLookup resolves caller-supplied path identifiers for a Turn or a -// Turn-scoped resource. Unparsable values are indistinguishable from missing ones. -func publicTurnLookup(tenantID, sessionID, turnID string) (sqlc.GetTurnParams, error) { - tenant, err := parseID(tenantID) - return sqlc.GetTurnParams{TenantID: tenant, SessionID: pgunit.PathID(sessionID), ID: pgunit.PathID(turnID)}, err -} diff --git a/services/core/internal/store/turns_test.go b/services/core/internal/store/turns_test.go index 297457606..4e85732e1 100644 --- a/services/core/internal/store/turns_test.go +++ b/services/core/internal/store/turns_test.go @@ -4,8 +4,6 @@ import ( "context" "encoding/json" "errors" - "reflect" - "sync" "testing" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -28,7 +26,7 @@ func TestCancellationStaysBoundToItsOriginalTurn(t *testing.T) { if err != nil || cancel.TurnID != first.TurnID { t.Fatalf("cancellation target: %+v, %v", cancel, err) } - pending, err := s.GetTurn(ctx, tenant, session.ID, first.TurnID) + pending, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, first.TurnID) if err != nil || pending.Status != sessions.TurnInProgress || pending.CancelRequestedAt.IsZero() || !pending.CompletedAt.IsZero() { t.Fatalf("running cancellation falsely completed: %+v, %v", pending, err) } @@ -40,18 +38,18 @@ func TestCancellationStaysBoundToItsOriginalTurn(t *testing.T) { t.Fatalf("cancellation retargeted: %+v, %v", retry, err) } } - queued, err := s.GetTurn(ctx, tenant, session.ID, next.TurnID) + queued, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, next.TurnID) if err != nil || queued.Status != sessions.TurnQueued || !queued.CancelRequestedAt.IsZero() { t.Fatalf("old cancellation affected later Turn: %+v, %v", queued, err) } if _, err := s.RequestCancel(ctx, tenant, session.ID, "cancel-queued"); err != nil { t.Fatal(err) } - stopped, err := s.GetTurn(ctx, tenant, session.ID, next.TurnID) + stopped, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, next.TurnID) if err != nil || stopped.Status != sessions.TurnCancelled || stopped.CompletedAt.IsZero() || !stopped.StartedAt.IsZero() { t.Fatalf("queued work did not stop: %+v, %v", stopped, err) } - if _, err := s.TransitionTurn(ctx, tenant, session.ID, next.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := transitionTurn(ctx, s, tenant, session.ID, next.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatalf("cancelled queued work was started: %v", err) } } @@ -75,64 +73,12 @@ func TestWaitingTurnRetainsInputsAndStartTime(t *testing.T) { if _, err := s.RequestCancel(ctx, tenant, session.ID, "cancel"); err != nil { t.Fatal(err) } - if _, err := s.TransitionTurn(ctx, tenant, session.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnWaiting, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := transitionTurn(ctx, s, tenant, session.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnWaiting, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatalf("cancelling Turn resumed: %v", err) } transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnWaiting, sessions.TurnCancelled) } -func TestTerminalOutcomeIsImmutableDuringConcurrentCallbacks(t *testing.T) { - s, pool := testStore(t) - other, _ := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - first := submitMessage(t, s, tenant, session.ID, "first") - transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - if _, err := s.RequestCancel(ctx, tenant, session.ID, "cancel"); err != nil { - t.Fatal(err) - } - var wg sync.WaitGroup - winners := make(chan sessions.Turn, 3) - errs := make(chan error, 3) - for _, status := range []string{sessions.TurnCompleted, sessions.TurnFailed, sessions.TurnCancelled} { - wg.Add(1) - go func() { - defer wg.Done() - outcome, _ := json.Marshal(map[string]string{"reported": status}) - turn, err := other.TransitionTurn(ctx, tenant, session.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: status, Outcome: outcome}) - if err == nil { - winners <- turn - } else { - errs <- err - } - }() - } - wg.Wait() - close(winners) - close(errs) - if len(winners) != 1 || len(errs) != 2 { - t.Fatalf("winners=%d errors=%d", len(winners), len(errs)) - } - for err := range errs { - if !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal(err) - } - } - winner := <-winners - if winner.CompletedAt.IsZero() || winner.CancelRequestedAt.IsZero() || winner.CompletedAt.Before(winner.StartedAt) { - t.Fatalf("terminal timestamps: %+v", winner) - } - if _, err := s.TransitionTurn(ctx, tenant, session.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"late":true}`)}); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatalf("late terminal callback accepted: %v", err) - } - pool.Close() - recovered, _ := testStore(t) - got, err := recovered.GetTurn(ctx, tenant, session.ID, first.TurnID) - if err != nil || !reflect.DeepEqual(got, winner) { - t.Fatalf("terminal outcome changed on restart: %+v, %v", got, err) - } -} - func TestTurnInputValidationHasNoSideEffects(t *testing.T) { s, _ := testStore(t) tenant, session := newTurnSession(t, s) @@ -157,11 +103,11 @@ func TestTurnInputValidationHasNoSideEffects(t *testing.T) { {ExpectedStatus: sessions.TurnCompleted, Status: sessions.TurnInProgress}, {ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress, Outcome: json.RawMessage(`{"premature":true}`)}, } { - if _, err := s.TransitionTurn(ctx, tenant, session.ID, first.TurnID, input); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := transitionTurn(ctx, s, tenant, session.ID, first.TurnID, input); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatalf("invalid transition accepted: %v", err) } } - got, err := s.GetTurn(ctx, tenant, session.ID, first.TurnID) + got, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, first.TurnID) if err != nil || got.Status != sessions.TurnQueued || !got.StartedAt.IsZero() { t.Fatalf("invalid transition changed Turn: %+v, %v", got, err) } diff --git a/services/core/internal/store/unified_model_configuration_http_test.go b/services/core/internal/store/unified_model_configuration_http_test.go index 9c7d4a57b..ad3b4ae6c 100644 --- a/services/core/internal/store/unified_model_configuration_http_test.go +++ b/services/core/internal/store/unified_model_configuration_http_test.go @@ -10,6 +10,7 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/api" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -113,7 +114,7 @@ func TestUnifiedModelConfigurationHTTP(t *testing.T) { t.Fatalf("wrong frozen selections: %#v", snapshot) } equalJSON(snapshot.HarnessConfig.Value, native) - secret, err := st.SessionModelExecution(t.Context(), tenant, id) + secret, err := store.SessionAdapter(st).SessionModelExecution(t.Context(), tenant, id) if err != nil || secret == nil || secret.APIKey != providerKey { t.Fatal("frozen provider changed", err) } diff --git a/services/core/internal/store/worker_capacity_test.go b/services/core/internal/store/worker_capacity_test.go index 596019ee1..327d776ef 100644 --- a/services/core/internal/store/worker_capacity_test.go +++ b/services/core/internal/store/worker_capacity_test.go @@ -7,6 +7,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) // Consume the actual controls so capacity rejection is exercised before any @@ -102,7 +103,7 @@ func TestWorkerDefersPreparationCapacityUntilCleanupReleasesSlot(t *testing.T) { } assertQueued := func() { t.Helper() - turn, err := h.s.GetTurn(t.Context(), h.tenant, blocked, turns[blocked]) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, blocked, turns[blocked]) if err != nil || turn.Status != sessions.TurnQueued { t.Fatal("capacity rejection failed queued work", turn, err) } @@ -135,7 +136,7 @@ func TestWorkerDefersPreparationCapacityUntilCleanupReleasesSlot(t *testing.T) { // still owns a slot and reports the same capacity rejection. h.write(turns[first], proto.TypeDone, proto.DonePayload{Content: "done"}) awaitDaemonRemoteCondition(t, t.Context(), 5*time.Second, "first completed Turn", func() bool { - turn, err := h.s.GetTurn(t.Context(), h.tenant, first, turns[first]) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, first, turns[first]) return err == nil && turn.Status == sessions.TurnCompleted }) observeCapacity() @@ -150,7 +151,7 @@ func TestWorkerDefersPreparationCapacityUntilCleanupReleasesSlot(t *testing.T) { } awaitDaemonRemoteCondition(t, t.Context(), 5*time.Second, "all five Turns completed", func() bool { for session, id := range turns { - turn, err := h.s.GetTurn(t.Context(), h.tenant, session, id) + turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, session, id) if err != nil || turn.Status != sessions.TurnCompleted { return false } diff --git a/services/core/internal/store/worker_input_race_test.go b/services/core/internal/store/worker_input_race_test.go index c24a1a287..c12653dc4 100644 --- a/services/core/internal/store/worker_input_race_test.go +++ b/services/core/internal/store/worker_input_race_test.go @@ -80,7 +80,7 @@ func TestWorkerInputReadSkipsConcurrentlyCancelledCandidate(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("input-read interleaving was not reached") } - turn, err := h.s.GetTurn(ctx, h.tenant, candidateSession, candidate.TurnID) + turn, err := store.SessionAdapter(h.s).GetTurn(ctx, h.tenant, candidateSession, candidate.TurnID) if err != nil || turn.Status != sessions.TurnCancelled { t.Fatal("candidate was not cancelled", err) } diff --git a/services/core/internal/store/worker_lease_loss_test.go b/services/core/internal/store/worker_lease_loss_test.go index 83ae5d02f..ef31f12a6 100644 --- a/services/core/internal/store/worker_lease_loss_test.go +++ b/services/core/internal/store/worker_lease_loss_test.go @@ -68,7 +68,7 @@ func TestWorkerLeaseLossLeavesUncertainWorkForSuccessor(t *testing.T) { case <-time.After(12 * time.Second): t.Fatal("worker ignored lease loss") } - active, err := h.s.GetTurn(t.Context(), h.tenant, h.session.ID, request.ID) + active, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, request.ID) if err != nil || active.Status != sessions.TurnInProgress { t.Fatal("lost owner persisted fallback completion", active, err) } @@ -78,11 +78,11 @@ func TestWorkerLeaseLossLeavesUncertainWorkForSuccessor(t *testing.T) { if err = successor.Run(stopped); err != context.Canceled { t.Fatal(err) } - active, err = h.s.GetTurn(t.Context(), h.tenant, h.session.ID, request.ID) + active, err = store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, request.ID) if err != nil || active.Status != sessions.TurnFailed { t.Fatal("successor did not reconcile", active, err) } - next, err := h.s.GetTurn(t.Context(), h.tenant, queued.ID, pending[0].TurnID) + next, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, queued.ID, pending[0].TurnID) if err != nil || next.Status != sessions.TurnQueued { t.Fatal("successor lost queued work", next, err) } diff --git a/services/core/tests/fixtures/items.go b/services/core/tests/fixtures/items.go index af1242f96..c0e1b6f9b 100644 --- a/services/core/tests/fixtures/items.go +++ b/services/core/tests/fixtures/items.go @@ -13,8 +13,8 @@ import ( ) // observeItems records the Turn's execution observations as the execution -// journal does, composing the journal procedure over a pooled Session -// transaction because the seeder holds no execution lease. +// journal does. The seeder holds no execution lease, so it composes the +// execution procedures over pooled Session transactions. func observeItems(ctx context.Context, pool *pgxpool.Pool, tenantID, sessionID, turn, status string) error { events := []sessions.ExecutionEvent{ {Kind: "delta", Payload: json.RawMessage(`{"item_id":"answer","delta":"partial answer"}`)}, @@ -31,6 +31,23 @@ func observeItems(ctx context.Context, pool *pgxpool.Pool, tenantID, sessionID, if err != nil { return err } + return inSession(ctx, pool, tenantID, sessionID, func(ctx context.Context, tx *sessionpg.SessionTx) error { + return sessions.AppendTurnEvents(ctx, tx, batch) + }) +} + +// transitionTurn moves the Turn as the execution owner does, running the +// transition procedure over a pooled Session transaction. +func transitionTurn(ctx context.Context, pool *pgxpool.Pool, tenantID, sessionID, turn string, transition sessions.TurnTransition) error { + return inSession(ctx, pool, tenantID, sessionID, func(ctx context.Context, tx *sessionpg.SessionTx) error { + _, err := sessions.TransitionTurn(ctx, tx, turn, transition) + return err + }) +} + +// inSession runs apply in a pooled transaction that holds the tenant's +// Session lock. +func inSession(ctx context.Context, pool *pgxpool.Pool, tenantID, sessionID string, apply func(context.Context, *sessionpg.SessionTx) error) error { tenant, err := pgunit.ParseID(tenantID) if err != nil { return err @@ -40,6 +57,6 @@ func observeItems(ctx context.Context, pool *pgxpool.Pool, tenantID, sessionID, return err } return sessionpg.WithSession(ctx, pgunit.NewPool(pool), tenant, session, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { - return sessions.AppendTurnEvents(ctx, sessionpg.BindSession(q, tenant, session), batch) + return apply(ctx, sessionpg.BindSession(q, tenant, session)) }) } diff --git a/services/core/tests/fixtures/main.go b/services/core/tests/fixtures/main.go index 8ef4db2b0..936afddfd 100644 --- a/services/core/tests/fixtures/main.go +++ b/services/core/tests/fixtures/main.go @@ -58,7 +58,7 @@ func seed() error { if err != nil { return err } - if _, err = s.TransitionTurn(ctx, f.Tenant, f.Session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { + if err = transitionTurn(ctx, pool, f.Tenant, f.Session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { return err } if err = observeItems(ctx, pool, f.Tenant, f.Session, receipt.TurnID, status); err != nil { @@ -66,7 +66,7 @@ func seed() error { } if status != sessions.TurnInProgress { outcome := json.RawMessage(`{"error":"SECRET engine log","done":{"metadata":{"agent_session_id":"PRIVATE"}}}`) - if _, err = s.TransitionTurn(ctx, f.Tenant, f.Session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: status, Outcome: outcome}); err != nil { + if err = transitionTurn(ctx, pool, f.Tenant, f.Session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: status, Outcome: outcome}); err != nil { return err } } From 25474e7a679b1b972f9b4b9291278fa7b5f04f55 Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 01:37:09 +0800 Subject: [PATCH 3/6] Stamp a sandbox reset with one clock reading (#435) --- .../core/internal/db/queries/sandbox_reset.sql | 16 +++++++++------- .../core/internal/db/sqlc/sandbox_reset.sql.go | 16 +++++++++------- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/services/core/internal/db/queries/sandbox_reset.sql b/services/core/internal/db/queries/sandbox_reset.sql index 58e819620..3fad61493 100644 --- a/services/core/internal/db/queries/sandbox_reset.sql +++ b/services/core/internal/db/queries/sandbox_reset.sql @@ -1,15 +1,17 @@ -- name: StartSandboxReset :exec +WITH clock AS MATERIALIZED (SELECT clock_timestamp() AS at) UPDATE runtime_deployment SET admission_paused = true, reset_clear = sqlc.arg(clear), - reset_requested_at = clock_timestamp(), + reset_requested_at = clock.at, reset_deadline_at = CASE WHEN sqlc.arg(clear)::text = 'auto' - THEN clock_timestamp() + make_interval(secs => sqlc.arg(deadline_seconds)::int) END, - reset_forced_at = CASE WHEN sqlc.arg(clear)::text = 'force' THEN clock_timestamp() END, - reset_audit = sqlc.arg(audit)::jsonb, updated_at = clock_timestamp() -WHERE singleton = true; + THEN clock.at + make_interval(secs => sqlc.arg(deadline_seconds)::int) END, + reset_forced_at = CASE WHEN sqlc.arg(clear)::text = 'force' THEN clock.at END, + reset_audit = sqlc.arg(audit)::jsonb, updated_at = clock.at +FROM clock WHERE singleton = true; -- name: ForceSandboxReset :exec -UPDATE runtime_deployment SET reset_clear = 'force', reset_forced_at = clock_timestamp(), updated_at = clock_timestamp() -WHERE singleton = true AND reset_clear = 'auto'; +WITH clock AS MATERIALIZED (SELECT clock_timestamp() AS at) +UPDATE runtime_deployment SET reset_clear = 'force', reset_forced_at = clock.at, updated_at = clock.at +FROM clock WHERE singleton = true AND reset_clear = 'auto'; -- name: CancelSandboxReset :exec UPDATE runtime_deployment SET admission_paused = false, reset_clear = NULL, diff --git a/services/core/internal/db/sqlc/sandbox_reset.sql.go b/services/core/internal/db/sqlc/sandbox_reset.sql.go index 6ded68955..49512c3f6 100644 --- a/services/core/internal/db/sqlc/sandbox_reset.sql.go +++ b/services/core/internal/db/sqlc/sandbox_reset.sql.go @@ -40,8 +40,9 @@ func (q *Queries) CompleteSandboxReset(ctx context.Context) error { } const forceSandboxReset = `-- name: ForceSandboxReset :exec -UPDATE runtime_deployment SET reset_clear = 'force', reset_forced_at = clock_timestamp(), updated_at = clock_timestamp() -WHERE singleton = true AND reset_clear = 'auto' +WITH clock AS MATERIALIZED (SELECT clock_timestamp() AS at) +UPDATE runtime_deployment SET reset_clear = 'force', reset_forced_at = clock.at, updated_at = clock.at +FROM clock WHERE singleton = true AND reset_clear = 'auto' ` func (q *Queries) ForceSandboxReset(ctx context.Context) error { @@ -228,13 +229,14 @@ func (q *Queries) SessionBlocksAutoReset(ctx context.Context, sessionID pgtype.U } const startSandboxReset = `-- name: StartSandboxReset :exec +WITH clock AS MATERIALIZED (SELECT clock_timestamp() AS at) UPDATE runtime_deployment SET admission_paused = true, reset_clear = $1, - reset_requested_at = clock_timestamp(), + reset_requested_at = clock.at, reset_deadline_at = CASE WHEN $1::text = 'auto' - THEN clock_timestamp() + make_interval(secs => $2::int) END, - reset_forced_at = CASE WHEN $1::text = 'force' THEN clock_timestamp() END, - reset_audit = $3::jsonb, updated_at = clock_timestamp() -WHERE singleton = true + THEN clock.at + make_interval(secs => $2::int) END, + reset_forced_at = CASE WHEN $1::text = 'force' THEN clock.at END, + reset_audit = $3::jsonb, updated_at = clock.at +FROM clock WHERE singleton = true ` type StartSandboxResetParams struct { From eed19fc241026ba2f46f2ce0d8d523cb538e2fdc Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 02:26:15 +0800 Subject: [PATCH 4/6] Move Session reads, deletion and archive into sessions (#438) * Move Session reads, deletion and archive into sessions Session reads, metadata updates, deletion and the public write audit move from store into sessions and sessionpg. GetSession, ListSessions and the stream snapshot each read from one snapshot, and the diagnostics read uses GetSession. sessionpg.LoadSessionActivity is the one Session activity read; store's Session creation and administrator summary call it. Session archive becomes deployment.ExecutionOperations.ArchiveSession and ArchiveResetSession over deploymentpg's WithSessionArchive: the Session lock, then the deployment's generation and reset checks in deployment, then cancellation through the sessions procedures, device revocation, cleanup request and audit in one leased transaction. The api archive route calls it directly. The Runtime observation resolver moves to deployment, and the Runtime history scope rule into runtimehistory.Service; both storeresolver packages are deleted. * Address Session reads review findings --- contracts/agents-api/runtime-observability.md | 2 +- .../agents-api/zh/runtime-observability.md | 4 +- services/core/IMPLEMENTATION.md | 8 +- services/core/cmd/server/http_routes_test.go | 1 + services/core/cmd/server/main.go | 22 +- .../internal/api/admin_session_archive.go | 4 +- .../api/admin_session_archive_test.go | 4 +- services/core/internal/api/dependencies.go | 3 +- .../core/internal/api/dependencies_test.go | 3 + .../core/internal/api/environment_input.go | 2 +- .../internal/api/environment_installation.go | 2 +- .../core/internal/api/errors_deployment.go | 4 +- services/core/internal/api/errors_test.go | 1 + services/core/internal/api/fakes_test.go | 74 ++-- services/core/internal/api/handler.go | 17 +- services/core/internal/api/handler_test.go | 2 +- services/core/internal/api/inputs.go | 2 +- .../native_classification_integration_test.go | 9 +- .../core/internal/api/resource_query_test.go | 10 +- .../internal/api/session_admission_test.go | 2 +- .../internal/api/session_creation_identity.go | 2 +- .../api/session_creation_stream_test.go | 4 +- .../core/internal/api/session_deletion.go | 2 +- .../core/internal/api/session_diagnostics.go | 5 +- .../api/session_diagnostics_deadline_test.go | 4 +- .../session_diagnostics_public_compat_test.go | 13 +- .../internal/api/session_diagnostics_test.go | 8 +- .../api/session_environment_http_test.go | 8 +- .../core/internal/api/session_metadata.go | 3 +- .../internal/api/session_semantics_test.go | 4 +- .../core/internal/api/session_write_audit.go | 3 +- .../internal/api/session_write_audit_test.go | 8 +- services/core/internal/api/stream.go | 2 +- services/core/internal/api/stream_test.go | 2 +- services/core/internal/api/turns.go | 4 +- services/core/internal/api/turns_test.go | 2 +- .../internal/api/validation_errors_test.go | 6 +- services/core/internal/deployment/errors.go | 3 + .../core/internal/deployment/fakes_test.go | 66 ++++ .../resolver.go => deployment/observation.go} | 42 +-- .../observation_test.go} | 93 +++-- .../internal/deployment/session_archive.go | 144 ++++++++ .../deployment/session_archive_test.go | 322 +++++++++++++++++ services/core/internal/deployment/storage.go | 46 +++ .../archive_cancellation_cleanup_test.go | 2 +- .../execution/deployment_fixture_test.go | 4 + .../core/internal/execution/dispatcher.go | 2 +- .../execution/environment_directory.go | 2 +- .../execution/environment_file_write.go | 2 +- .../internal/execution/prepared_dispatch.go | 2 +- .../core/internal/execution/sandbox_reset.go | 5 +- .../internal/execution/session_archive.go | 13 - services/core/internal/execution/worker.go | 2 +- .../core/internal/execution/worker_device.go | 2 +- .../postgres/deploymentpg/allocations.go | 44 +-- .../postgres/deploymentpg/session_archive.go | 85 +++++ .../persistence/postgres/deploymentpg/tx.go | 38 +- .../persistence/postgres/sessionpg/binding.go | 6 + .../postgres/sessionpg/session_reads.go | 325 ++++++++++++++++++ .../postgres/sessionpg/session_writes.go | 86 +++++ .../postgres/sessionpg/turn_reads.go | 6 +- .../core/internal/runtimehistory/service.go | 47 ++- .../internal/runtimehistory/service_test.go | 80 ++++- .../runtimehistory/storeresolver/resolver.go | 48 --- .../storeresolver/resolver_test.go | 74 ---- .../providers/configuration_flow_test.go | 9 +- .../core/internal/sessions/artifacts_test.go | 15 + .../core/internal/sessions/devices_test.go | 20 ++ .../sessions/execution_configuration.go | 44 +++ services/core/internal/sessions/reader.go | 3 +- services/core/internal/sessions/reads.go | 49 ++- services/core/internal/sessions/service.go | 1 + services/core/internal/sessions/session.go | 117 +++++++ .../core/internal/sessions/session_test.go | 102 ++++++ .../internal/sessions/transaction_test.go | 14 + .../internal/store/admin_delete_audit_test.go | 8 +- .../internal/store/admin_session_archive.go | 161 --------- .../store/admin_session_archive_race_test.go | 8 +- .../store/admin_session_archive_test.go | 44 ++- .../admin_session_archive_worker_http_test.go | 27 +- services/core/internal/store/admin_summary.go | 2 +- .../agent_execution_defaults_http_test.go | 7 +- .../store/archive_cancellation_test.go | 5 +- .../archived_cancellation_migration_test.go | 2 +- .../internal/store/command_output_test.go | 12 +- .../creation_stream_settlement_public_test.go | 7 +- .../deployment_model_providers_http_test.go | 4 +- .../store/environment_admission_test.go | 4 +- .../store/environment_claim_worker_test.go | 2 +- .../environment_connection_events_test.go | 2 +- .../store/environment_connections_test.go | 4 +- .../internal/store/environment_device_test.go | 2 +- .../store/environment_directory_test.go | 2 +- .../environment_executor_credentials_test.go | 4 +- .../environment_executor_management_test.go | 2 +- ...onment_file_write_semantics_public_test.go | 4 +- .../store/environment_file_writes_test.go | 4 +- .../store/environment_initial_input_test.go | 16 +- .../store/environment_initialization_test.go | 2 +- .../store/environment_input_activity_test.go | 30 +- .../store/environment_input_claim_test.go | 8 +- .../store/environment_input_expiry_test.go | 2 +- .../environment_input_settlement_test.go | 2 +- .../store/environment_installation_test.go | 2 +- .../store/environment_mcp_public_test.go | 2 +- .../environment_skill_references_test.go | 2 +- .../internal/store/environment_work_test.go | 2 +- .../store/environment_worker_helpers_test.go | 2 +- .../store/environment_worker_scan_test.go | 3 +- .../core/internal/store/environments_test.go | 8 +- .../core/internal/store/execution_test.go | 29 +- .../store/executor_principals_test.go | 4 +- services/core/internal/store/export_test.go | 3 + .../internal/store/function_execution_test.go | 2 +- .../store/function_images_native_test.go | 2 +- .../internal/store/function_inputs_test.go | 2 +- .../store/function_item_events_test.go | 4 +- .../store/function_public_native_test.go | 2 +- .../internal/store/function_state_test.go | 18 +- .../internal/store/harness_onboarding_test.go | 6 +- ...sted_initialization_failure_public_test.go | 20 +- .../internal/store/item_serialization_test.go | 6 +- .../internal/store/list_query_public_test.go | 2 +- .../store/local_environment_devices_test.go | 2 +- .../local_environment_file_write_test.go | 2 +- .../store/local_environment_worker_test.go | 4 +- .../store/mcode_public_native_test.go | 2 +- .../store/message_image_admission_test.go | 2 +- .../store/message_images_native_test.go | 2 +- .../store/model_protocol_native_test.go | 2 +- .../store/native_public_execution_test.go | 2 +- .../store/prepared_dispatch_failure_test.go | 2 +- .../internal/store/prepared_dispatch_test.go | 2 +- .../internal/store/public_execution_test.go | 10 +- .../store/public_handler_fixture_test.go | 18 +- .../store/runtime_allocations_test.go | 2 +- .../store/runtime_compute_lifecycle_test.go | 4 +- .../runtime_configuration_cleanup_test.go | 8 +- .../internal/store/runtime_connection_test.go | 5 +- .../store/runtime_creation_settlement_test.go | 6 +- .../internal/store/runtime_deployment_test.go | 4 +- .../runtime_enrollment_connection_test.go | 2 +- .../internal/store/runtime_enrollment_test.go | 2 +- .../runtime_environment_terminal_test.go | 10 +- .../store/runtime_lifecycle_nodes_test.go | 2 +- .../internal/store/runtime_lifecycle_test.go | 4 +- .../store/runtime_node_lifecycle_test.go | 6 +- services/core/internal/store/runtime_nodes.go | 13 - .../core/internal/store/runtime_nodes_test.go | 14 +- .../store/runtime_observation_scan_test.go | 5 +- .../internal/store/runtime_pending_test.go | 6 +- .../core/internal/store/runtime_scan_test.go | 4 +- .../internal/store/runtime_suspension_test.go | 4 +- .../store/runtime_worker_recovery_test.go | 2 +- .../store/sandbox_deployment_switch_test.go | 4 +- .../sandbox_deployment_switch_worker_test.go | 3 +- .../core/internal/store/sandbox_reset_test.go | 26 +- .../store/sandbox_specification_store_test.go | 2 +- services/core/internal/store/scheduling.go | 101 ------ .../store/self_hosted_cancel_public_test.go | 6 +- .../store/self_hosted_initial_public_test.go | 2 +- .../store/session_agent_filter_test.go | 10 +- .../internal/store/session_artifacts_test.go | 6 +- .../store/session_configuration_test.go | 6 +- .../store/session_creation_stream_test.go | 16 +- .../internal/store/session_creator_test.go | 8 +- .../core/internal/store/session_deletion.go | 64 ---- .../store/session_deletion_execution_test.go | 18 +- .../internal/store/session_deletion_test.go | 45 +-- .../internal/store/session_diagnostics.go | 87 ----- .../session_diagnostics_cancellation_test.go | 8 +- .../store/session_diagnostics_test.go | 34 +- .../session_environment_snapshot_test.go | 4 +- .../core/internal/store/session_events.go | 67 ---- .../internal/store/session_events_test.go | 26 +- .../store/session_execution_configuration.go | 70 +--- .../session_execution_configuration_test.go | 22 +- .../store/session_initial_input_test.go | 8 +- .../store/session_initial_public_test.go | 4 +- .../core/internal/store/session_metadata.go | 46 --- .../internal/store/session_metadata_test.go | 16 +- .../session_reference_retry_public_test.go | 2 +- .../internal/store/session_transaction.go | 20 +- .../internal/store/session_write_audit.go | 22 -- .../store/session_write_audit_test.go | 20 +- services/core/internal/store/sessions.go | 64 +--- services/core/internal/store/sessions_test.go | 12 +- .../store/structured_output_native_test.go | 2 +- .../store/subagent_identities_test.go | 6 +- .../store/subagent_native_outputs_test.go | 2 +- .../internal/store/subagent_resources_test.go | 2 +- .../store/template_composition_public_test.go | 2 +- .../template_null_selection_public_test.go | 2 +- services/core/internal/store/token_usage.go | 30 -- .../store/token_usage_integration_test.go | 16 +- .../internal/store/tool_policy_native_test.go | 8 +- .../internal/store/tool_search_native_test.go | 2 +- .../core/internal/store/turn_inputs_test.go | 2 +- .../unified_model_configuration_http_test.go | 2 +- .../store/whitespace_input_public_test.go | 2 +- .../store/worker_preparation_failure_test.go | 5 +- 201 files changed, 2245 insertions(+), 1559 deletions(-) rename services/core/internal/{runtimeobs/storeresolver/resolver.go => deployment/observation.go} (79%) rename services/core/internal/{runtimeobs/storeresolver/resolver_test.go => deployment/observation_test.go} (70%) create mode 100644 services/core/internal/deployment/session_archive.go create mode 100644 services/core/internal/deployment/session_archive_test.go delete mode 100644 services/core/internal/execution/session_archive.go create mode 100644 services/core/internal/persistence/postgres/deploymentpg/session_archive.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/session_reads.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/session_writes.go delete mode 100644 services/core/internal/runtimehistory/storeresolver/resolver.go delete mode 100644 services/core/internal/runtimehistory/storeresolver/resolver_test.go create mode 100644 services/core/internal/sessions/execution_configuration.go create mode 100644 services/core/internal/sessions/session_test.go delete mode 100644 services/core/internal/store/admin_session_archive.go delete mode 100644 services/core/internal/store/runtime_nodes.go delete mode 100644 services/core/internal/store/session_deletion.go delete mode 100644 services/core/internal/store/session_diagnostics.go delete mode 100644 services/core/internal/store/session_events.go delete mode 100644 services/core/internal/store/session_metadata.go delete mode 100644 services/core/internal/store/session_write_audit.go delete mode 100644 services/core/internal/store/token_usage.go diff --git a/contracts/agents-api/runtime-observability.md b/contracts/agents-api/runtime-observability.md index f1a961d79..226fc2ca8 100644 --- a/contracts/agents-api/runtime-observability.md +++ b/contracts/agents-api/runtime-observability.md @@ -16,7 +16,7 @@ self-hosted: tenant_id -> session_id -> environment_id -> device_id + connection none: tenant_id -> session_id (no Session-owned Runtime instance) ``` -The resolver (`internal/runtimeobs/storeresolver`) reads the Session, its Environment, the current allocation and the Session's measured usage from the store. A Session, daemon connection, process, container and native Harness Session are different identities and never stand in for one another. +The resolver (`services/core/internal/deployment/observation.go`) reads the Session, its Environment, the current allocation and the Session's measured usage from the database. A Session, daemon connection, process, container and native Harness Session are different identities and never stand in for one another. Managed Docker, microsandbox and E2B allocations are observed. `none` and `self_hosted` Sessions are `unsupported`; Core never attributes shared host statistics to an `environment:none` Session. diff --git a/contracts/agents-api/zh/runtime-observability.md b/contracts/agents-api/zh/runtime-observability.md index a6d894a45..8e24b9a56 100644 --- a/contracts/agents-api/zh/runtime-observability.md +++ b/contracts/agents-api/zh/runtime-observability.md @@ -1,7 +1,7 @@ --- title: "运行时可观测性" source: contracts/agents-api/runtime-observability.md -source_hash: 5da1279a81a6dcb3a85661dbba942c8937f871ae351ed80550db65b2a59156db +source_hash: f79a077350118f5f9bb3f4e8874242af3cb716e92f5e090f6652b001ae5780a1 --- 这是面向贡献者的契约，规定 Core 如何观测 Runtime 并保留其历史。路由和响应字段见 [Runtime telemetry API](runtime-observability-api.md)。代码位于 `services/core/internal/runtimeobs`（解析、源、采样器和导出）、`internal/runtimehistory`（历史查询和 PostgreSQL 存储）以及 `internal/runtimeobs/otlpexporter`。 @@ -18,7 +18,7 @@ self-hosted: tenant_id -> session_id -> environment_id -> device_id + connection none: tenant_id -> session_id (no Session-owned Runtime instance) ``` -解析器（`internal/runtimeobs/storeresolver`）从存储中读取 Session、其 Environment、当前分配以及 Session 的实测使用量。Session、守护进程连接、进程、容器和原生 Harness Session 是不同身份，彼此绝不能替代。 +解析器（`services/core/internal/deployment/observation.go`）从数据库中读取 Session、其 Environment、当前分配以及 Session 的实测使用量。Session、守护进程连接、进程、容器和原生 Harness Session 是不同身份，彼此绝不能替代。 托管 Docker、microsandbox 和 E2B 分配均会被观测。`none` 和 `self_hosted` Session 为 `unsupported`；Core 绝不会将共享主机统计信息归属于 `environment:none` Session。 diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 2254cb2de..338a330bc 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, `LoadRequiredActions` lists the calls a Turn waits on, and `TransitionTurn` moves a root Turn's status as its current status allows and settles a Turn that ends. `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 and Turn operation loads the Turn 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` or `TransitionTurn`, in that transaction; `LoadEnvironment`, which reads the tenant's Environment of a live Session; and `CreateSessionEnvironment`, which stores the Environment of a Session being created. The other `sessionpg` exports serve `store` until the Session operations leave it. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; `placementpg`, which loads the placement facts and reserves placements, sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. +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` or `TransitionTurn`, 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; among them `LoadSessionActivity`, the one Session activity read, adds the activity projection to the Sessions that `store`'s Session creation and administrator summary return. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; `placementpg`, which loads the placement facts and reserves placements, sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. `store` is transitional. `store.New` builds a pooled Store, and `store.NewExecution(s, lease)` builds the execution writer on a lease it borrows. An execution-only operation on a pooled Store fails with `store.ErrExecutionAuthority`. New adapters do not copy that check: their execution repositories require a `*pgunit.Lease` at construction, their public repositories expose no execution operation, and the check goes away with `store`. @@ -33,10 +33,10 @@ Domain owners, each with its PostgreSQL adapter under `internal/persistence/post - `environmenttemplates` (`templatepg`): Environment Templates, their validation and default network, their sealed setup, initial files, Skills and Plugins, and the resolved Template that Session creation composes into its Environment. - `modelconfiguration` (`modelconfigurationpg`): each Harness's deployment default model configuration and its last-use observations. - `skills` (`skillpg`): Skills and their immutable versions: archive checks, the default and latest pointers, version selection and deletion, and each version's sealed archive. Session creation freezes selected versions inside its `store` transaction with the `skills` rules. -- `deployment` and its subpackage `deployment/placement` (`deploymentpg`, `placementpg`): the sandbox deployment and its nodes: provider configuration and the sealed credential, the specification and retained generations, setup, update and switch, node enrollment, identity and authentication, generation configuration, capacity, presence, status, host history, reset, and the counts of nodes and sandboxes bound to the public address; and hosted runtime allocations: reservation with the dedicated device, compute ownership and settlement, observation diagnostics, activity, compute phases, wake receipts and cleanup, with the reads that schedule, discover and authorize them. It interprets Sandbox Provider declarations through the `providers.Registry` it is given, whose lookups return typed errors. `cmd/server` builds that registry and calls it only to build a direct Provider and to discover a Provider's configuration; it takes each setup's mode and declared operations from the `Setup` that `deployment` returns. `deployment/placement` owns hosted admission and placement: `cmd/server` builds one `placement.Rules` from the registry and the public URL, both fixed while Core runs, and gives it to `deployment.Service` and to `store`; its pure decisions admit a hosted Session, choose its node, and admit an allocation's reserved node and a restore on it, and its errors keep one status, code and message through `writeStoreError` and `writeDeploymentError`. `placementpg` loads the facts those decisions read and applies them on the caller's transaction-bound queries; it has no Store or transaction runner. The only other reader is `store`: Session creation decides admission and placement with those rules and `placementpg` inside its own transaction. Deployment changes and allocation writes run through `deployment.ExecutionOperations` on `deploymentpg.NewExecution(lease, …)`, which the Worker receives as `execution.Owner.Deployment`. Each allocation write locks the owning Session, decides in `deployment`, applies in `deploymentpg` and prunes the Session's change journal in the same leased transaction; cleanup settles the Session through the `sessions` procedures on a `sessionpg.SessionTx` bound to that transaction. The Worker reads the deployment, prepares a selection's setup and records live-compute activity through the pooled `deployment.Service` in `execution.Dispatcher.Deployment`, and lists the Sessions a reset still has to archive, schedules node lifecycles and reads allocations through the pooled `deployment.Reader` in `execution.Dispatcher.DeploymentReader`; a pooled read carries no lease, and the leased write that follows it rechecks the allocation's owner. Node management and reads use the pooled `deploymentpg.Store`, which `cmd/server` also reads the owner epoch from and runs the host-history sampler on. Provider calls run outside transactions, and the final transaction rechecks the expected generation. The reset's Session archive stays in `store`. +- `deployment` and its subpackage `deployment/placement` (`deploymentpg`, `placementpg`): the sandbox deployment and its nodes: provider configuration and the sealed credential, the specification and retained generations, setup, update and switch, node enrollment, identity and authentication, generation configuration, capacity, presence, status, host history, reset, and the counts of nodes and sandboxes bound to the public address; and hosted runtime allocations: reservation with the dedicated device, compute ownership and settlement, observation diagnostics, activity, compute phases, wake receipts and cleanup, with the reads that schedule, discover and authorize them. It interprets Sandbox Provider declarations through the `providers.Registry` it is given, whose lookups return typed errors. `cmd/server` builds that registry and calls it only to build a direct Provider and to discover a Provider's configuration; it takes each setup's mode and declared operations from the `Setup` that `deployment` returns. `deployment/placement` owns hosted admission and placement: `cmd/server` builds one `placement.Rules` from the registry and the public URL, both fixed while Core runs, and gives it to `deployment.Service` and to `store`; its pure decisions admit a hosted Session, choose its node, and admit an allocation's reserved node and a restore on it, and its errors keep one status, code and message through `writeStoreError` and `writeDeploymentError`. `placementpg` loads the facts those decisions read and applies them on the caller's transaction-bound queries; it has no Store or transaction runner. The only other reader is `store`: Session creation decides admission and placement with those rules and `placementpg` inside its own transaction. Deployment changes and allocation writes run through `deployment.ExecutionOperations` on `deploymentpg.NewExecution(lease, …)`, which the Worker receives as `execution.Owner.Deployment`. Each allocation write locks the owning Session, decides in `deployment`, applies in `deploymentpg` and prunes the Session's change journal in the same leased transaction; cleanup settles the Session through the `sessions` procedures on a `sessionpg.SessionTx` bound to that transaction. The Worker reads the deployment, prepares a selection's setup and records live-compute activity through the pooled `deployment.Service` in `execution.Dispatcher.Deployment`, and lists the Sessions a reset still has to archive, schedules node lifecycles and reads allocations through the pooled `deployment.Reader` in `execution.Dispatcher.DeploymentReader`; a pooled read carries no lease, and the leased write that follows it rechecks the allocation's owner. Node management and reads use the pooled `deploymentpg.Store`, which `cmd/server` also reads the owner epoch from and runs the host-history sampler on. Provider calls run outside transactions, and the final transaction rechecks the expected generation. Session archive crosses the Session and the deployment, so it is a `deployment.ExecutionOperations` operation: `ArchiveSession`, which the administrator's archive route calls through `api.Execution.SessionArchive`, and `ArchiveResetSession`, which a reset calls, lock the Session through `deploymentpg`'s `WithSessionArchive`, check the deployment's generation and the running reset in `deployment`, then expire the Environment and cancel its work through the `sessions` procedures, revoke the allocation's device and request its cleanup, and record the audit entry in one leased transaction. `deployment.ObservationResolver` resolves a Session's Runtime observation target for `runtimeobs` from `sessions.SessionReader` and `deployment.Reader`. - `coremetrics` (`coremetricspg`): the Core metrics PostgreSQL holds: the root Turn queue counts, the root Turn history, read from one read-only snapshot, and the database size. `cmd/server`'s Core metrics source adds them to the process, pool, Worker and daemon registry measurements. -- `runtimehistory` (`runtimehistory/postgresreader`, outside this directory): Runtime history samples: the periodic export, scoped reads and retention, which also prunes node-host samples. -- `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 Turns (their reads and the execution work scan), the model provider a Session froze, which `sessionpg.Store` opens with the credential key, 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, which include a Session's execution binding and the execution device list, their creation, revocation, heartbeats and Runtime enrollment, and the archived cancellation receipt the daemon gateway reads), the Sessions whose Runtime the administrator observes across Projects, 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. Turn transitions, execution completion, which alone publishes or discards a Turn's staged Artifacts, the start of Artifact capture, 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`. +- `runtimehistory` (`runtimehistory/postgresreader`, outside this directory): Runtime history samples: the periodic export, scoped reads and retention, which also prunes node-host samples. `runtimehistory.Service` scopes a read to the Session's hosted Environment through `sessions.EnvironmentReader`. +- `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 Sessions (their reads, the change journal and stream snapshot, diagnostics, the frozen execution configuration, measured usage and archive state, metadata updates, deletion and the public write audit), Turns (their reads and the execution work scan), the model provider a Session froze, which `sessionpg.Store` opens with the credential key, 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, which include a Session's execution binding and the execution device list, their creation, revocation, heartbeats and Runtime enrollment, and the archived cancellation receipt the daemon gateway reads), the Sessions whose Runtime the administrator observes across Projects, 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. Turn transitions, execution completion, which alone publishes or discards a Turn's staged Artifacts, the start of Artifact capture, function calls and their application receipts, the Turn execution journal, Environment initialization, connection observations and their reconciliation, Session device binding, and file-write reservation and settlement run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. ## Request handling diff --git a/services/core/cmd/server/http_routes_test.go b/services/core/cmd/server/http_routes_test.go index 998788235..05585fd38 100644 --- a/services/core/cmd/server/http_routes_test.go +++ b/services/core/cmd/server/http_routes_test.go @@ -97,6 +97,7 @@ func daemonComposition(t testing.TB) http.Handler { Agents: struct{ api.Agents }{}, AgentsReader: struct{ api.AgentsReader }{}, EnvironmentTemplates: struct{ api.EnvironmentTemplates }{}, EnvironmentTemplatesReader: struct{ api.EnvironmentTemplatesReader }{}, Sessions: struct{ api.Sessions }{}, + SessionsReader: struct{ api.SessionsReader }{}, SessionCreation: struct{ api.SessionCreation }{}, SessionEvents: struct{ api.SessionEvents }{}, Turns: struct{ api.Turns }{}, diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index 6cf3a08ad..3cf6979dc 100644 --- a/services/core/cmd/server/main.go +++ b/services/core/cmd/server/main.go @@ -61,9 +61,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeenrollment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimehistory" - historystoreresolver "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimehistory/storeresolver" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" - observationstoreresolver "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs/storeresolver" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/providers" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/skills" @@ -230,7 +228,7 @@ func run() error { managed = managedNodes.runtime observationSources[managed.InstallationID] = managedNodes.setup } - observationResolver, err := observationstoreresolver.NewResolver(executionStore, deploymentStore) + observationResolver, err := deployment.NewObservationResolver(sessionStore, deploymentStore) if err != nil { return err } @@ -274,11 +272,7 @@ func run() error { if err := api.ValidateCredentialSeparation(ctx, keyAdmin, projectStore); err != nil { return err } - historyResolver, err := historystoreresolver.NewResolver(sessionStore) - if err != nil { - return err - } - historyService, err := runtimehistory.NewService(historyResolver, history.Reader) + historyService, err := runtimehistory.NewService(sessionStore, history.Reader) if err != nil { return err } @@ -307,6 +301,7 @@ func run() error { nativeInstaller = &api.NativeInstaller{Version: buildRevision, Catalog: catalog} } } + var deploymentExecution *deployment.ExecutionOperations if registry != nil { dispatcher := &execution.Dispatcher{Store: executionStore, Registry: registry, Credentials: vaultService, Observer: modelConfigurationStore, Deployment: deploymentService, DeploymentReader: deploymentStore, @@ -317,7 +312,7 @@ func run() error { if err != nil { return err } - deploymentExecution, err := deployment.NewExecutionOperations(deploymentService, deploymentpg.NewExecution(lease, credentialKey)) + deploymentExecution, err = deployment.NewExecutionOperations(deploymentService, deploymentpg.NewExecution(lease, credentialKey)) if err != nil { return errors.Join(err, lease.Close(ctx)) } @@ -403,15 +398,16 @@ func run() error { EnvironmentTemplates: environmentTemplates, EnvironmentTemplatesReader: templateStore, Files: fileService, FilesReader: fileStore, Agents: agentService, AgentsReader: agentStore, - Sessions: executionStore, + Sessions: sessionService, + SessionsReader: sessionStore, SessionCreation: executionStore, - SessionEvents: executionStore, + SessionEvents: sessionStore, Turns: sessionStore, Items: sessionStore, Subagents: sessionStore, Artifacts: sessionService, ArtifactsReader: sessionStore, - SessionAdmin: executionStore, + SessionAdmin: sessionStore, Environments: sessionService, EnvironmentsReader: sessionStore, ExecutorConnections: executorConnections{sessions: sessionStore, registry: registry}, Admin: executionStore, AdminRuntimeTargets: sessionStore, AdminAudit: auditStore, WriteAudit: auditStore, Metrics: metrics, RuntimeObservations: observationService, RuntimeHistory: historyService, @@ -421,7 +417,7 @@ func run() error { ExecutorURL: executorURL, SessionAdmission: worker, InputAdmission: worker, - SessionArchive: worker, + SessionArchive: deploymentExecution, Workspaces: worker, NativeInstaller: nativeInstaller, } diff --git a/services/core/internal/api/admin_session_archive.go b/services/core/internal/api/admin_session_archive.go index a17268514..bb50213fc 100644 --- a/services/core/internal/api/admin_session_archive.go +++ b/services/core/internal/api/admin_session_archive.go @@ -11,7 +11,7 @@ import ( // SessionArchive archives a managed Session through the execution owner; // SessionAdmin reads its archive state. type SessionArchive interface { - ArchiveManagedSession(context.Context, string, string, uint64) (sessions.ManagedArchive, error) + ArchiveSession(context.Context, string, string, uint64) (sessions.ManagedArchive, error) } type AdminSessionArchiveRequest struct { @@ -44,7 +44,7 @@ func (h *Handler) adminArchiveSession(w http.ResponseWriter, r *http.Request) { writeStoreError(w, r, sessions.ErrEnvironmentUnavailable) return } - result, err := h.Execution.SessionArchive.ArchiveManagedSession(r.Context(), tenantID(r), chi.URLParam(r, "session_id"), input.ExpectedGeneration) + result, err := h.Execution.SessionArchive.ArchiveSession(r.Context(), tenantID(r), chi.URLParam(r, "session_id"), input.ExpectedGeneration) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/admin_session_archive_test.go b/services/core/internal/api/admin_session_archive_test.go index 6ef7550ff..4d12f04a4 100644 --- a/services/core/internal/api/admin_session_archive_test.go +++ b/services/core/internal/api/admin_session_archive_test.go @@ -19,7 +19,7 @@ type archiveManagementFixture struct { err error } -func (s *archiveManagementFixture) ArchiveManagedSession(ctx context.Context, tenant, session string, generation uint64) (sessions.ManagedArchive, error) { +func (s *archiveManagementFixture) ArchiveSession(ctx context.Context, tenant, session string, generation uint64) (sessions.ManagedArchive, error) { s.calls++ s.tenant, s.session, s.generation = tenant, session, generation source, ok := adminaudit.FromContext(ctx) @@ -39,7 +39,7 @@ func TestAdminSessionArchiveAuthorityAndValidation(t *testing.T) { deps, fakes := managementFakes(t, key) fixture := &archiveManagementFixture{} deps.Execution = fakes.execution() - fakes.sessionArchive.archiveManagedSession = fixture.ArchiveManagedSession + fakes.sessionArchive.archiveSession = fixture.ArchiveSession fakes.sessionAdmin.getManagedSessionArchive = fixture.GetManagedSessionArchive h := newTestHandler(t, deps) path := "/core/v1/projects/" + managementProjectID + "/sessions/11111111-1111-4111-8111-111111111111/archive" diff --git a/services/core/internal/api/dependencies.go b/services/core/internal/api/dependencies.go index c810fa9c6..9c7a5e090 100644 --- a/services/core/internal/api/dependencies.go +++ b/services/core/internal/api/dependencies.go @@ -43,6 +43,7 @@ type Dependencies struct { Agents Agents AgentsReader AgentsReader Sessions Sessions + SessionsReader SessionsReader SessionCreation SessionCreation SessionEvents SessionEvents Turns Turns @@ -133,7 +134,7 @@ func (d Dependencies) validate() error { field{"EnvironmentTemplates", d.EnvironmentTemplates}, field{"EnvironmentTemplatesReader", d.EnvironmentTemplatesReader}, field{"Skills", d.Skills}, field{"SkillsReader", d.SkillsReader}, field{"Agents", d.Agents}, field{"AgentsReader", d.AgentsReader}, - field{"Sessions", d.Sessions}, + field{"Sessions", d.Sessions}, field{"SessionsReader", d.SessionsReader}, field{"SessionCreation", d.SessionCreation}, field{"SessionEvents", d.SessionEvents}, field{"Turns", d.Turns}, diff --git a/services/core/internal/api/dependencies_test.go b/services/core/internal/api/dependencies_test.go index 698629487..38ce94de5 100644 --- a/services/core/internal/api/dependencies_test.go +++ b/services/core/internal/api/dependencies_test.go @@ -29,6 +29,7 @@ type testFakes struct { agents *fakeAgents agentsReader *fakeAgentsReader sessions *fakeSessions + sessionsReader *fakeSessionsReader sessionCreation *fakeSessionCreation sessionEvents *fakeSessionEvents turns *fakeTurns @@ -77,6 +78,7 @@ func testDependencies(t testing.TB) (Dependencies, *testFakes) { skills: &fakeSkills{t: t}, skillsReader: &fakeSkillsReader{t: t}, agents: &fakeAgents{t: t}, agentsReader: &fakeAgentsReader{t: t}, sessions: &fakeSessions{t: t}, + sessionsReader: &fakeSessionsReader{t: t}, sessionCreation: &fakeSessionCreation{t: t}, sessionEvents: &fakeSessionEvents{t: t}, turns: &fakeTurns{t: t}, @@ -104,6 +106,7 @@ func testDependencies(t testing.TB) (Dependencies, *testFakes) { Skills: f.skills, SkillsReader: f.skillsReader, Agents: f.agents, AgentsReader: f.agentsReader, Sessions: f.sessions, + SessionsReader: f.sessionsReader, SessionCreation: f.sessionCreation, SessionEvents: f.sessionEvents, Turns: f.turns, diff --git a/services/core/internal/api/environment_input.go b/services/core/internal/api/environment_input.go index 69c4e1ae9..fed8b537c 100644 --- a/services/core/internal/api/environment_input.go +++ b/services/core/internal/api/environment_input.go @@ -9,7 +9,7 @@ import ( ) func (h *Handler) setEnvironmentInputWriteDeadline(w http.ResponseWriter, r *http.Request, sessionID string) error { - session, err := h.Sessions.GetSession(r.Context(), tenantID(r), sessionID) + session, err := h.SessionsReader.GetSession(r.Context(), tenantID(r), sessionID) if err != nil { return err } diff --git a/services/core/internal/api/environment_installation.go b/services/core/internal/api/environment_installation.go index 80412dcd8..c9074702f 100644 --- a/services/core/internal/api/environment_installation.go +++ b/services/core/internal/api/environment_installation.go @@ -118,7 +118,7 @@ func (h *Handler) prepareNativeInstallation(w http.ResponseWriter, r *http.Reque writeSessionsError(w, r, err) return } - session, err := h.Sessions.GetSession(r.Context(), claim.Principal.TenantID, environment.SessionID) + session, err := h.SessionsReader.GetSession(r.Context(), claim.Principal.TenantID, environment.SessionID) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/errors_deployment.go b/services/core/internal/api/errors_deployment.go index 352b39a2b..eeb8a212b 100644 --- a/services/core/internal/api/errors_deployment.go +++ b/services/core/internal/api/errors_deployment.go @@ -22,8 +22,6 @@ func writeDeploymentError(w http.ResponseWriter, r *http.Request, err error) { writeError(w, http.StatusServiceUnavailable, "execution_unavailable", "Execution is not available on this service.") case errors.Is(err, deployment.ErrNotFound): writeError(w, http.StatusNotFound, "not_found_error", "Resource not found.") - case errors.Is(err, deployment.ErrInvalidInput): - writeError(w, http.StatusBadRequest, "invalid_request", invalidInputMessage) default: if writeAuditSourceError(w, r, err) || writeTextValueError(w, r, err) || writeCredentialUnavailableError(w, r, err) { return @@ -106,6 +104,8 @@ func writeSandboxError(w http.ResponseWriter, err error) bool { writeError(w, http.StatusServiceUnavailable, "sandbox_nodes_preparing", "Sandbox nodes are preparing the requested Runtime.") case errors.Is(err, placement.ErrNodeUnavailable): writeError(w, http.StatusServiceUnavailable, "runtime_node_unavailable", "The selected sandbox node is unavailable or has no capacity.") + case errors.Is(err, deployment.ErrInvalidInput): + writeError(w, http.StatusBadRequest, "invalid_request", invalidInputMessage) default: return false } diff --git a/services/core/internal/api/errors_test.go b/services/core/internal/api/errors_test.go index 6d88ecab7..d7bd2cb31 100644 --- a/services/core/internal/api/errors_test.go +++ b/services/core/internal/api/errors_test.go @@ -189,6 +189,7 @@ func TestSharedPersistenceErrors(t *testing.T) { {writeAuditError, adminaudit.ErrInvalidQuery, 400, invalid}, {storeError, fmt.Errorf("write: %w", textvalue.ErrUnstorable), 400, unstorableTextMessage}, {writeAuditError, textvalue.ErrUnstorable, 400, unstorableTextMessage}, + {storeError, deployment.ErrInvalidInput, 400, invalid}, {storeError, credentialcrypto.ErrUnavailable, 503, "credential_storage_unavailable"}, {writeAuditError, errors.New("canary"), 500, "internal_error"}, } { diff --git a/services/core/internal/api/fakes_test.go b/services/core/internal/api/fakes_test.go index bd21b0a51..cf1565aee 100644 --- a/services/core/internal/api/fakes_test.go +++ b/services/core/internal/api/fakes_test.go @@ -739,19 +739,11 @@ func (f *fakeRuntimeObservations) ObserveSessions(a0 context.Context, a1 []runti type fakeSessionAdmin struct { t testing.TB - getSessionDiagnosticsSnapshot func(context.Context, string, string) (sessions.Session, error) getTurnDiagnosticsSnapshot func(context.Context, string, string, string) (sessions.TurnDiagnosticsSnapshot, error) getSessionExecutionConfiguration func(context.Context, string, string) (v1.SessionExecutionConfiguration, error) getManagedSessionArchive func(context.Context, string, string) (sessions.ManagedArchive, error) } -func (f *fakeSessionAdmin) GetSessionDiagnosticsSnapshot(a0 context.Context, a1 string, a2 string) (sessions.Session, error) { - if f.getSessionDiagnosticsSnapshot == nil { - unexpectedCall(f.t, "GetSessionDiagnosticsSnapshot") - } - return f.getSessionDiagnosticsSnapshot(a0, a1, a2) -} - func (f *fakeSessionAdmin) GetTurnDiagnosticsSnapshot(a0 context.Context, a1 string, a2 string, a3 string) (sessions.TurnDiagnosticsSnapshot, error) { if f.getTurnDiagnosticsSnapshot == nil { unexpectedCall(f.t, "GetTurnDiagnosticsSnapshot") @@ -794,15 +786,15 @@ func (f *fakeSessionAdmission) CreateSessionStream(a0 context.Context, a1 string } type fakeSessionArchive struct { - t testing.TB - archiveManagedSession func(context.Context, string, string, uint64) (sessions.ManagedArchive, error) + t testing.TB + archiveSession func(context.Context, string, string, uint64) (sessions.ManagedArchive, error) } -func (f *fakeSessionArchive) ArchiveManagedSession(a0 context.Context, a1 string, a2 string, a3 uint64) (sessions.ManagedArchive, error) { - if f.archiveManagedSession == nil { - unexpectedCall(f.t, "ArchiveManagedSession") +func (f *fakeSessionArchive) ArchiveSession(a0 context.Context, a1 string, a2 string, a3 uint64) (sessions.ManagedArchive, error) { + if f.archiveSession == nil { + unexpectedCall(f.t, "ArchiveSession") } - return f.archiveManagedSession(a0, a1, a2, a3) + return f.archiveSession(a0, a1, a2, a3) } type fakeSessionCreation struct { @@ -863,46 +855,50 @@ func (f *fakeSessionEvents) SessionStreamSnapshot(a0 context.Context, a1 string, type fakeSessions struct { t testing.TB - getSession func(context.Context, string, string) (sessions.Session, error) - listSessions func(context.Context, string, string, int, bool, *string) (sessions.Page, error) - updateSessionMetadata func(context.Context, string, string, map[string]string) (sessions.Session, error) - deleteSession func(context.Context, string, string) error - auditSessionOperation func(context.Context, string, string, string) error -} - -func (f *fakeSessions) GetSession(a0 context.Context, a1 string, a2 string) (sessions.Session, error) { - if f.getSession == nil { - unexpectedCall(f.t, "GetSession") - } - return f.getSession(a0, a1, a2) -} - -func (f *fakeSessions) ListSessions(a0 context.Context, a1 string, a2 string, a3 int, a4 bool, a5 *string) (sessions.Page, error) { - if f.listSessions == nil { - unexpectedCall(f.t, "ListSessions") - } - return f.listSessions(a0, a1, a2, a3, a4, a5) + updateSessionMetadata func(context.Context, sessions.UpdateSessionMetadataCommand) (sessions.Session, error) + deleteSession func(context.Context, sessions.DeleteSessionCommand) error + auditSessionOperation func(context.Context, sessions.AuditSessionOperationCommand) error } -func (f *fakeSessions) UpdateSessionMetadata(a0 context.Context, a1 string, a2 string, a3 map[string]string) (sessions.Session, error) { +func (f *fakeSessions) UpdateSessionMetadata(a0 context.Context, a1 sessions.UpdateSessionMetadataCommand) (sessions.Session, error) { if f.updateSessionMetadata == nil { unexpectedCall(f.t, "UpdateSessionMetadata") } - return f.updateSessionMetadata(a0, a1, a2, a3) + return f.updateSessionMetadata(a0, a1) } -func (f *fakeSessions) DeleteSession(a0 context.Context, a1 string, a2 string) error { +func (f *fakeSessions) DeleteSession(a0 context.Context, a1 sessions.DeleteSessionCommand) error { if f.deleteSession == nil { unexpectedCall(f.t, "DeleteSession") } - return f.deleteSession(a0, a1, a2) + return f.deleteSession(a0, a1) } -func (f *fakeSessions) AuditSessionOperation(a0 context.Context, a1 string, a2 string, a3 string) error { +func (f *fakeSessions) AuditSessionOperation(a0 context.Context, a1 sessions.AuditSessionOperationCommand) error { if f.auditSessionOperation == nil { unexpectedCall(f.t, "AuditSessionOperation") } - return f.auditSessionOperation(a0, a1, a2, a3) + return f.auditSessionOperation(a0, a1) +} + +type fakeSessionsReader struct { + t testing.TB + getSession func(context.Context, string, string) (sessions.Session, error) + listSessions func(context.Context, string, string, int, bool, *string) (sessions.Page, error) +} + +func (f *fakeSessionsReader) GetSession(a0 context.Context, a1 string, a2 string) (sessions.Session, error) { + if f.getSession == nil { + unexpectedCall(f.t, "GetSession") + } + return f.getSession(a0, a1, a2) +} + +func (f *fakeSessionsReader) ListSessions(a0 context.Context, a1 string, a2 string, a3 int, a4 bool, a5 *string) (sessions.Page, error) { + if f.listSessions == nil { + unexpectedCall(f.t, "ListSessions") + } + return f.listSessions(a0, a1, a2, a3, a4, a5) } type fakeSkills struct { diff --git a/services/core/internal/api/handler.go b/services/core/internal/api/handler.go index 77746f5b2..154b0a491 100644 --- a/services/core/internal/api/handler.go +++ b/services/core/internal/api/handler.go @@ -34,14 +34,17 @@ type SessionAdmission interface { CreateSessionStream(context.Context, string, sessions.CreateSession) (sessions.Creation, error) } -// Sessions reads, updates and deletes Sessions, and records their public write -// audit. +// Sessions updates and deletes Sessions, and records their public write audit. type Sessions interface { + UpdateSessionMetadata(context.Context, sessions.UpdateSessionMetadataCommand) (sessions.Session, error) + DeleteSession(context.Context, sessions.DeleteSessionCommand) error + AuditSessionOperation(context.Context, sessions.AuditSessionOperationCommand) error +} + +// SessionsReader reads Sessions. +type SessionsReader interface { GetSession(context.Context, string, string) (sessions.Session, error) ListSessions(context.Context, string, string, int, bool, *string) (sessions.Page, error) - UpdateSessionMetadata(context.Context, string, string, map[string]string) (sessions.Session, error) - DeleteSession(context.Context, string, string) error - AuditSessionOperation(context.Context, string, string, string) error } // routes builds the router. HEAD runs the GET route without a body after the @@ -302,7 +305,7 @@ func (h *Handler) createSession(w http.ResponseWriter, r *http.Request) { // @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")) + session, err := h.SessionsReader.GetSession(r.Context(), tenantID(r), chi.URLParam(r, "session_id")) if err != nil { writeStoreError(w, r, err) return @@ -349,7 +352,7 @@ func (h *Handler) listSessions(w http.ResponseWriter, r *http.Request) { if values, present := r.URL.Query()["agent_id"]; present { agentID = &values[0] } - page, err := h.Sessions.ListSessions(r.Context(), tenantID(r), options.after, options.limit, options.ascending, agentID) + page, err := h.SessionsReader.ListSessions(r.Context(), tenantID(r), options.after, options.limit, options.ascending, agentID) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/handler_test.go b/services/core/internal/api/handler_test.go index cdcd8d501..0a0ab0536 100644 --- a/services/core/internal/api/handler_test.go +++ b/services/core/internal/api/handler_test.go @@ -52,7 +52,7 @@ func (s *recordingStore) CreateSession(_ context.Context, tenant string, input s // record answers Session creation, reads and listing from s. func (s *recordingStore) record(f *testFakes) { f.sessionCreation.createSession, f.sessionCreation.findSessionCreation = s.CreateSession, s.FindSessionCreation - f.sessions.getSession, f.sessions.listSessions = s.GetSession, s.ListSessions + f.sessionsReader.getSession, f.sessionsReader.listSessions = s.GetSession, s.ListSessions } // testHandler serves strict fakes for a fresh tenant whose caller diff --git a/services/core/internal/api/inputs.go b/services/core/internal/api/inputs.go index 94f9af3e4..aa6f6ad37 100644 --- a/services/core/internal/api/inputs.go +++ b/services/core/internal/api/inputs.go @@ -58,7 +58,7 @@ func (h *Handler) createEvents(w http.ResponseWriter, r *http.Request) { if request.Events != nil && len(request.Events) == 0 { // Empty batches have no execution identity to reserve or replay. // Authorize the resource even when no executor is configured. - if _, err := h.Sessions.GetSession(r.Context(), tenantID(r), chi.URLParam(r, "session_id")); err != nil { + if _, err := h.SessionsReader.GetSession(r.Context(), tenantID(r), chi.URLParam(r, "session_id")); err != nil { writeStoreError(w, r, err) return } diff --git a/services/core/internal/api/native_classification_integration_test.go b/services/core/internal/api/native_classification_integration_test.go index 5b0a4530a..9ce5ea206 100644 --- a/services/core/internal/api/native_classification_integration_test.go +++ b/services/core/internal/api/native_classification_integration_test.go @@ -9,13 +9,16 @@ 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/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" ) func TestNativeClassificationPostgresRoundTripAndPublicPrivacy(t *testing.T) { s, pool := diagnosticDatabase(t) - h, _, tenant := adminTestHandler(t, databaseSessionReads(s, pool)) + h, _, tenant := adminTestHandler(t, databaseSessionReads(pool)) + reader := sessionpg.New(pgunit.NewPool(pool), nil) for _, code := range []string{"authentication_error", "connection_failed", "secret-canary"} { t.Run(code, func(t *testing.T) { session, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: identity.Subject{Kind: "service_account", ID: "native-classification"}, Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"id":"agent_root","model":"test"},"environment":{"type":"none"}}`)}) @@ -34,7 +37,7 @@ func TestNativeClassificationPostgresRoundTripAndPublicPrivacy(t *testing.T) { t.Fatal(err) } transitionTurn(t, pool, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: outcome}) - snap, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + snap, err := reader.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil { t.Fatal(err) } @@ -65,7 +68,7 @@ func TestNativeClassificationPostgresRoundTripAndPublicPrivacy(t *testing.T) { } before[i] = append([]byte(nil), w.Body.Bytes()...) } - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + events, err := reader.ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/api/resource_query_test.go b/services/core/internal/api/resource_query_test.go index 9e4cf7408..bf5dde96c 100644 --- a/services/core/internal/api/resource_query_test.go +++ b/services/core/internal/api/resource_query_test.go @@ -54,12 +54,12 @@ func (s *missingResourceStore) GetSession(_ context.Context, tenant, _ string) ( return sessions.Session{}, s.missing(tenant) } -func (s *missingResourceStore) DeleteSession(_ context.Context, tenant, _ string) error { - return s.missing(tenant) +func (s *missingResourceStore) DeleteSession(_ context.Context, command sessions.DeleteSessionCommand) error { + return s.missing(command.TenantID) } -func (s *missingResourceStore) UpdateSessionMetadata(_ context.Context, tenant, _ string, _ map[string]string) (sessions.Session, error) { - return sessions.Session{}, s.missing(tenant) +func (s *missingResourceStore) UpdateSessionMetadata(_ context.Context, command sessions.UpdateSessionMetadataCommand) (sessions.Session, error) { + return sessions.Session{}, s.missing(command.TenantID) } // Environment Template operations report a missing Template with their @@ -84,7 +84,7 @@ func (s *missingResourceStore) DeleteEnvironmentTemplate(_ context.Context, comm // wire serves the Agent, Session and Environment template lookups from s. func (s *missingResourceStore) wire(_ *Dependencies, f *testFakes) { f.agentsReader.getAgent, f.agents.delete, f.agents.update = s.GetAgent, s.DeleteAgent, s.UpdateAgent - f.sessions.getSession, f.sessions.deleteSession, f.sessions.updateSessionMetadata = s.GetSession, s.DeleteSession, s.UpdateSessionMetadata + f.sessionsReader.getSession, f.sessions.deleteSession, f.sessions.updateSessionMetadata = s.GetSession, s.DeleteSession, s.UpdateSessionMetadata f.environmentTemplatesReader.get, f.environmentTemplates.update, f.environmentTemplates.delete = s.GetEnvironmentTemplate, s.UpdateEnvironmentTemplate, s.DeleteEnvironmentTemplate } diff --git a/services/core/internal/api/session_admission_test.go b/services/core/internal/api/session_admission_test.go index 49914f9f4..aedd04a7b 100644 --- a/services/core/internal/api/session_admission_test.go +++ b/services/core/internal/api/session_admission_test.go @@ -15,7 +15,7 @@ import ( // testHandler serves, so any access fails the test. func forbidSessionAccess(_ *Dependencies, f *testFakes) { f.sessionCreation.createSession, f.sessionCreation.findSessionCreation = nil, nil - f.sessions.getSession, f.sessions.listSessions = nil, nil + f.sessionsReader.getSession, f.sessionsReader.listSessions = nil, nil f.modelProviders.resolve = nil } diff --git a/services/core/internal/api/session_creation_identity.go b/services/core/internal/api/session_creation_identity.go index 22a0bb248..c466a6d1c 100644 --- a/services/core/internal/api/session_creation_identity.go +++ b/services/core/internal/api/session_creation_identity.go @@ -60,7 +60,7 @@ func (h *Handler) recoverSessionCreation(w http.ResponseWriter, r *http.Request, // Recorded-intent lookup finds an existing creation, which sends no events. h.respondSessionCreationStream(w, r, result) } else { - session, err := h.Sessions.GetSession(r.Context(), tenantID(r), result.Session.ID) + session, err := h.SessionsReader.GetSession(r.Context(), tenantID(r), result.Session.ID) if err != nil { writeStoreError(w, r, err) } else if h.auditSessionOperation(w, r, session.ID, "create") { diff --git a/services/core/internal/api/session_creation_stream_test.go b/services/core/internal/api/session_creation_stream_test.go index bb6dfc839..a94119d7d 100644 --- a/services/core/internal/api/session_creation_stream_test.go +++ b/services/core/internal/api/session_creation_stream_test.go @@ -102,7 +102,7 @@ func (f *creationStreamFixture) FindSessionCreation(context.Context, string, str } // AuditSessionOperation accepts the audit of a replayed creation. -func (f *creationStreamFixture) AuditSessionOperation(context.Context, string, string, string) error { +func (f *creationStreamFixture) AuditSessionOperation(context.Context, sessions.AuditSessionOperationCommand) error { return nil } @@ -146,7 +146,7 @@ func newCreationStreamHarness(t *testing.T) *creationStreamHarness { OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential("key"), TenantID: tenant, }).ResolveAPIKey - fakes.sessions.getSession, fakes.sessions.auditSessionOperation = fixture.GetSession, fixture.AuditSessionOperation + fakes.sessionsReader.getSession, fakes.sessions.auditSessionOperation = fixture.GetSession, fixture.AuditSessionOperation fakes.sessionCreation.findSessionCreation = fixture.FindSessionCreation fakes.sessionEvents.sessionEventCursor, fakes.sessionEvents.sessionStreamSnapshot, fakes.sessionEvents.listSessionEvents = fixture.SessionEventCursor, fixture.SessionStreamSnapshot, fixture.ListSessionEvents fakes.modelProviders.resolve = noDeploymentModelProvider diff --git a/services/core/internal/api/session_deletion.go b/services/core/internal/api/session_deletion.go index ce3917733..7905c0c81 100644 --- a/services/core/internal/api/session_deletion.go +++ b/services/core/internal/api/session_deletion.go @@ -35,7 +35,7 @@ func (h *Handler) deleteSession(w http.ResponseWriter, r *http.Request) { writeStoreError(w, r, sessions.ErrNotFound) return } - if err := h.Sessions.DeleteSession(r.Context(), tenantID(r), id); err != nil { + if err := h.Sessions.DeleteSession(r.Context(), sessions.DeleteSessionCommand{TenantID: tenantID(r), SessionID: id}); err != nil { writeStoreError(w, r, err) return } diff --git a/services/core/internal/api/session_diagnostics.go b/services/core/internal/api/session_diagnostics.go index 138534f00..a98f57312 100644 --- a/services/core/internal/api/session_diagnostics.go +++ b/services/core/internal/api/session_diagnostics.go @@ -46,10 +46,9 @@ type TurnDiagnostics struct { ItemsTruncated bool `json:"items_truncated"` } -// SessionAdmin serves the administrator's per-Session reads: diagnostics +// SessionAdmin serves the administrator's per-Session reads: Turn diagnostics // snapshots, the execution configuration and the managed archive state. type SessionAdmin interface { - GetSessionDiagnosticsSnapshot(context.Context, string, string) (sessions.Session, error) GetTurnDiagnosticsSnapshot(context.Context, string, string, string) (sessions.TurnDiagnosticsSnapshot, error) GetSessionExecutionConfiguration(context.Context, string, string) (v1.SessionExecutionConfiguration, error) GetManagedSessionArchive(context.Context, string, string) (sessions.ManagedArchive, error) @@ -68,7 +67,7 @@ type SessionAdmin interface { func (h *Handler) getSessionDiagnostics(w http.ResponseWriter, r *http.Request) { ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) defer cancel() - session, err := h.SessionAdmin.GetSessionDiagnosticsSnapshot(ctx, tenantID(r), chi.URLParam(r, "session_id")) + session, err := h.SessionsReader.GetSession(ctx, tenantID(r), chi.URLParam(r, "session_id")) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/session_diagnostics_deadline_test.go b/services/core/internal/api/session_diagnostics_deadline_test.go index e8fd9445e..523840f37 100644 --- a/services/core/internal/api/session_diagnostics_deadline_test.go +++ b/services/core/internal/api/session_diagnostics_deadline_test.go @@ -17,9 +17,9 @@ type diagnosticDeadlineStore struct { observed context.Context } -func (s *diagnosticDeadlineStore) GetSessionDiagnosticsSnapshot(ctx context.Context, tenant, session string) (sessions.Session, error) { +func (s *diagnosticDeadlineStore) GetSession(ctx context.Context, tenant, session string) (sessions.Session, error) { s.observed = ctx - return s.diagnosticSnapshotStore.GetSessionDiagnosticsSnapshot(ctx, tenant, session) + return s.diagnosticSnapshotStore.GetSession(ctx, tenant, session) } func (s *diagnosticDeadlineStore) GetTurnDiagnosticsSnapshot(ctx context.Context, tenant, session, turn string) (sessions.TurnDiagnosticsSnapshot, error) { s.observed = ctx diff --git a/services/core/internal/api/session_diagnostics_public_compat_test.go b/services/core/internal/api/session_diagnostics_public_compat_test.go index 6910ec7db..c555210b3 100644 --- a/services/core/internal/api/session_diagnostics_public_compat_test.go +++ b/services/core/internal/api/session_diagnostics_public_compat_test.go @@ -27,7 +27,7 @@ func TestDiagnosticPublicCompatibility(t *testing.T) { key := callerBinding() deps, fakes := testDependencies(t) fakes.projectsReader.resolveAPIKey = projectKeys(t, key).ResolveAPIKey - databaseSessionReads(s, pool)(&deps, fakes) + databaseSessionReads(pool)(&deps, fakes) h := newTestHandler(t, deps) session, err := s.CreateSession(t.Context(), key.TenantID, sessions.CreateSession{Creator: identity.Subject{Kind: "service_account", ID: "compat-test"}, Engine: "codex", IdempotencyKey: uuid.NewString(), Configuration: json.RawMessage(`{"agent":{"id":"agent_root","model":"test"},"environment":{"type":"none"}}`)}) if err != nil { @@ -64,13 +64,12 @@ func diagnosticRequest(handler http.Handler, path, token string) *httptest.Respo return w } -// databaseSessionReads serves Session, Turn and diagnostic reads from s, and -// Item reads from the Session adapter on pool. -func databaseSessionReads(s *store.Store, pool *pgxpool.Pool) func(*Dependencies, *testFakes) { +// databaseSessionReads serves Session, Turn, diagnostic and Item reads from +// the Session adapter on pool. +func databaseSessionReads(pool *pgxpool.Pool) func(*Dependencies, *testFakes) { return func(d *Dependencies, _ *testFakes) { - d.Sessions, d.SessionAdmin = s, s - reads := sessionpg.New(pgunit.NewPool(pool), nil) - d.Items, d.Turns = reads, reads + reader := sessionpg.New(pgunit.NewPool(pool), nil) + d.SessionsReader, d.SessionAdmin, d.Items, d.Turns = reader, reader, reader, reader } } diff --git a/services/core/internal/api/session_diagnostics_test.go b/services/core/internal/api/session_diagnostics_test.go index 8fc433ae0..448a959df 100644 --- a/services/core/internal/api/session_diagnostics_test.go +++ b/services/core/internal/api/session_diagnostics_test.go @@ -14,7 +14,7 @@ import ( func TestDiagnosticsCoreHandlerDatabaseBoundary(t *testing.T) { s, pool := diagnosticDatabase(t) - h, _, tenant := adminTestHandler(t, databaseSessionReads(s, pool)) + h, _, tenant := adminTestHandler(t, databaseSessionReads(pool)) session, err := s.CreateSession(t.Context(), tenant, sessions.CreateSession{Creator: identity.Subject{Kind: "service_account", ID: "diagnostic-test"}, Engine: "codex", IdempotencyKey: "diagnostics", Configuration: json.RawMessage(`{"agent":{"id":"agent_root","model":"test"},"environment":{"type":"none"}}`)}) if err != nil { t.Fatal(err) @@ -61,7 +61,7 @@ type diagnosticSnapshotStore struct { session sessions.Session } -func (s diagnosticSnapshotStore) GetSessionDiagnosticsSnapshot(context.Context, string, string) (sessions.Session, error) { +func (s diagnosticSnapshotStore) GetSession(context.Context, string, string) (sessions.Session, error) { return s.session, nil } func (s diagnosticSnapshotStore) GetTurnDiagnosticsSnapshot(context.Context, string, string, string) (sessions.TurnDiagnosticsSnapshot, error) { @@ -70,14 +70,14 @@ func (s diagnosticSnapshotStore) GetTurnDiagnosticsSnapshot(context.Context, str // diagnosticSnapshots answers Core diagnostic reads. type diagnosticSnapshots interface { - GetSessionDiagnosticsSnapshot(context.Context, string, string) (sessions.Session, error) + GetSession(context.Context, string, string) (sessions.Session, error) GetTurnDiagnosticsSnapshot(context.Context, string, string, string) (sessions.TurnDiagnosticsSnapshot, error) } // serveDiagnostics answers Session and Turn diagnostic reads from source. func serveDiagnostics(source diagnosticSnapshots) func(*Dependencies, *testFakes) { return func(_ *Dependencies, f *testFakes) { - f.sessionAdmin.getSessionDiagnosticsSnapshot, f.sessionAdmin.getTurnDiagnosticsSnapshot = source.GetSessionDiagnosticsSnapshot, source.GetTurnDiagnosticsSnapshot + f.sessionsReader.getSession, f.sessionAdmin.getTurnDiagnosticsSnapshot = source.GetSession, source.GetTurnDiagnosticsSnapshot } } diff --git a/services/core/internal/api/session_environment_http_test.go b/services/core/internal/api/session_environment_http_test.go index 9d93c79cd..d607345c2 100644 --- a/services/core/internal/api/session_environment_http_test.go +++ b/services/core/internal/api/session_environment_http_test.go @@ -25,9 +25,9 @@ func (f *environmentHTTPFixture) ListSessions(_ context.Context, tenant, _ strin return sessions.Page{Sessions: []sessions.Session{f.session}}, nil } -func (f *environmentHTTPFixture) UpdateSessionMetadata(ctx context.Context, tenant, session string, metadata map[string]string) (sessions.Session, error) { - value, err := f.GetSession(ctx, tenant, session) - value.Metadata = metadata +func (f *environmentHTTPFixture) UpdateSessionMetadata(ctx context.Context, command sessions.UpdateSessionMetadataCommand) (sessions.Session, error) { + value, err := f.GetSession(ctx, command.TenantID, command.SessionID) + value.Metadata = command.Metadata return value, err } @@ -48,7 +48,7 @@ func TestSelfHostedSessionHTTPReadListMetadataAndLiveStream(t *testing.T) { TokenSHA256: runtimedevice.HashCredential("key"), TenantID: session.TenantID, }).ResolveAPIKey fixture.serve(fakes) - fakes.sessions.listSessions, fakes.sessions.updateSessionMetadata = fixture.ListSessions, fixture.UpdateSessionMetadata + fakes.sessionsReader.listSessions, fakes.sessions.updateSessionMetadata = fixture.ListSessions, fixture.UpdateSessionMetadata // Self-hosted Sessions report the executor URL of the enabled Execution. deps.Execution = fakes.execution() deps.Execution.ExecutorURL = environmentOrigin diff --git a/services/core/internal/api/session_metadata.go b/services/core/internal/api/session_metadata.go index e4c334fb5..6cec3c02b 100644 --- a/services/core/internal/api/session_metadata.go +++ b/services/core/internal/api/session_metadata.go @@ -10,6 +10,7 @@ import ( "slices" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/metadata" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/go-chi/chi/v5" ) @@ -59,7 +60,7 @@ func (h *Handler) updateSession(w http.ResponseWriter, r *http.Request) { } return } - session, err := h.Sessions.UpdateSessionMetadata(r.Context(), tenantID(r), chi.URLParam(r, "session_id"), values) + session, err := h.Sessions.UpdateSessionMetadata(r.Context(), sessions.UpdateSessionMetadataCommand{TenantID: tenantID(r), SessionID: chi.URLParam(r, "session_id"), Metadata: values}) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/session_semantics_test.go b/services/core/internal/api/session_semantics_test.go index fe5809db0..2cb5b1ceb 100644 --- a/services/core/internal/api/session_semantics_test.go +++ b/services/core/internal/api/session_semantics_test.go @@ -27,7 +27,7 @@ func (s *emptyEventSessionStore) GetSession(_ context.Context, tenant, id string return sessions.Session{ID: id, TenantID: tenant, Configuration: json.RawMessage(`{"environment":{"type":"none"}}`)}, nil } -func (s *emptyEventSessionStore) AuditSessionOperation(context.Context, string, string, string) error { +func (s *emptyEventSessionStore) AuditSessionOperation(context.Context, sessions.AuditSessionOperationCommand) error { return nil } @@ -37,7 +37,7 @@ func TestEmptyEventBatchAuthorizesWithoutExecutionEffects(t *testing.T) { recorder := &inputRecorder{} sessions := &emptyEventSessionStore{} h, _, tenant := testHandler(t, func(d *Dependencies, f *testFakes) { - f.sessions.getSession, f.sessions.auditSessionOperation = sessions.GetSession, sessions.AuditSessionOperation + f.sessionsReader.getSession, f.sessions.auditSessionOperation = sessions.GetSession, sessions.AuditSessionOperation if executor { recorder.admit(d, f) } diff --git a/services/core/internal/api/session_write_audit.go b/services/core/internal/api/session_write_audit.go index 52cb4e3be..afc117383 100644 --- a/services/core/internal/api/session_write_audit.go +++ b/services/core/internal/api/session_write_audit.go @@ -3,6 +3,7 @@ package api import ( "net/http" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/writeaudit" ) @@ -10,7 +11,7 @@ func (h *Handler) auditSessionOperation(w http.ResponseWriter, r *http.Request, if _, ok := writeaudit.FromContext(r.Context()); !ok { return true } - if err := h.Sessions.AuditSessionOperation(r.Context(), tenantID(r), sessionID, action); err != nil { + if err := h.Sessions.AuditSessionOperation(r.Context(), sessions.AuditSessionOperationCommand{TenantID: tenantID(r), SessionID: sessionID, Action: action}); err != nil { writeStoreError(w, r, err) return false } diff --git a/services/core/internal/api/session_write_audit_test.go b/services/core/internal/api/session_write_audit_test.go index cd0987761..a341ce362 100644 --- a/services/core/internal/api/session_write_audit_test.go +++ b/services/core/internal/api/session_write_audit_test.go @@ -27,12 +27,12 @@ func (f *auditedSessionFixture) FindSessionCreation(context.Context, string, str return sessions.Creation{Session: f.session}, nil } -func (f *auditedSessionFixture) AuditSessionOperation(ctx context.Context, tenant, session, action string) error { - if tenant != f.session.TenantID || session != f.session.ID { +func (f *auditedSessionFixture) AuditSessionOperation(ctx context.Context, command sessions.AuditSessionOperationCommand) error { + if command.TenantID != f.session.TenantID || command.SessionID != f.session.ID { return sessions.ErrNotFound } f.source, _ = writeaudit.FromContext(ctx) - f.actions = append(f.actions, action) + f.actions = append(f.actions, command.Action) return f.err } @@ -48,7 +48,7 @@ func TestSessionAuditOnlyRoutesFailClosed(t *testing.T) { deps, fakes := testDependencies(t) fakes.sessions.auditSessionOperation = f.AuditSessionOperation if route != "stream-replay" { - fakes.sessions.getSession = f.GetSession + fakes.sessionsReader.getSession = f.GetSession } if route != "empty-events" { fakes.sessionCreation.findSessionCreation = f.FindSessionCreation diff --git a/services/core/internal/api/stream.go b/services/core/internal/api/stream.go index 1168ee433..ba7cff7f0 100644 --- a/services/core/internal/api/stream.go +++ b/services/core/internal/api/stream.go @@ -37,7 +37,7 @@ type SessionEvents interface { // @Router /agents/sessions/{session_id}/events [get] func (h *Handler) streamEvents(w http.ResponseWriter, r *http.Request) { id, tenant := chi.URLParam(r, "session_id"), tenantID(r) - session, err := h.Sessions.GetSession(r.Context(), tenant, id) + session, err := h.SessionsReader.GetSession(r.Context(), tenant, id) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/stream_test.go b/services/core/internal/api/stream_test.go index a82b707f3..1e376e0e6 100644 --- a/services/core/internal/api/stream_test.go +++ b/services/core/internal/api/stream_test.go @@ -61,7 +61,7 @@ func (f *streamFixture) ListSessionEvents(_ context.Context, _, _ string, cursor // serve answers Session reads and the event stream from f. func (f *streamFixture) serve(fakes *testFakes) { - fakes.sessions.getSession = f.GetSession + fakes.sessionsReader.getSession = f.GetSession fakes.sessionEvents.sessionEventCursor, fakes.sessionEvents.sessionStreamSnapshot, fakes.sessionEvents.listSessionEvents = f.SessionEventCursor, f.SessionStreamSnapshot, f.ListSessionEvents } diff --git a/services/core/internal/api/turns.go b/services/core/internal/api/turns.go index 7da89680a..cfd2668e6 100644 --- a/services/core/internal/api/turns.go +++ b/services/core/internal/api/turns.go @@ -36,7 +36,7 @@ func (h *Handler) getTurn(w http.ResponseWriter, r *http.Request) { writeStoreError(w, r, err) return } - session, err := h.Sessions.GetSession(r.Context(), tenantID(r), sessionID) + session, err := h.SessionsReader.GetSession(r.Context(), tenantID(r), sessionID) if err != nil { writeStoreError(w, r, err) return @@ -68,7 +68,7 @@ func (h *Handler) listTurns(w http.ResponseWriter, r *http.Request) { return } sessionID := chi.URLParam(r, "session_id") - session, err := h.Sessions.GetSession(r.Context(), tenantID(r), sessionID) + session, err := h.SessionsReader.GetSession(r.Context(), tenantID(r), sessionID) if err != nil { writeStoreError(w, r, err) return diff --git a/services/core/internal/api/turns_test.go b/services/core/internal/api/turns_test.go index 47acb988a..cbf738706 100644 --- a/services/core/internal/api/turns_test.go +++ b/services/core/internal/api/turns_test.go @@ -37,7 +37,7 @@ func (s *turnReadStore) ListTurns(_ context.Context, tenant, session, cursor str func TestTurnRoutesUseAuthenticatedScopeAndSafeProjection(t *testing.T) { s := &turnReadStore{session: sessions.Session{Configuration: json.RawMessage(`{"agent":{"id":"agent_snapshot"}}`)}, turn: sessions.Turn{ID: "turn", SessionID: "session", Status: sessions.TurnFailed, CreatedAt: time.Unix(1700000000, 999), Outcome: json.RawMessage(`{"error":"Bearer SECRET","done":{"metadata":{"password":"SECRET"}}}`)}} h, _, tenant := testHandler(t, func(_ *Dependencies, f *testFakes) { - f.sessions.getSession, f.turns.getTurn, f.turns.listTurns = s.GetSession, s.GetTurn, s.ListTurns + f.sessionsReader.getSession, f.turns.getTurn, f.turns.listTurns = s.GetSession, s.GetTurn, s.ListTurns }) request := func(path string) *httptest.ResponseRecorder { r := httptest.NewRequest(http.MethodGet, path, nil) diff --git a/services/core/internal/api/validation_errors_test.go b/services/core/internal/api/validation_errors_test.go index 2820e6e24..9c05b9458 100644 --- a/services/core/internal/api/validation_errors_test.go +++ b/services/core/internal/api/validation_errors_test.go @@ -40,9 +40,9 @@ func (s *validationStore) CreateVault(_ context.Context, input vaults.CreateVaul return vaults.Vault{ID: uuid.NewString(), TenantID: input.TenantID, Name: input.Name, Metadata: input.Metadata}, nil } -func (s *validationStore) UpdateSessionMetadata(_ context.Context, tenant, id string, metadata map[string]string) (sessions.Session, error) { +func (s *validationStore) UpdateSessionMetadata(_ context.Context, command sessions.UpdateSessionMetadataCommand) (sessions.Session, error) { s.writes++ - return sessions.Session{ID: id, TenantID: tenant, Metadata: metadata, Configuration: json.RawMessage(`{"agent":{"id":"agent_validation","model":"validation-model"},"environment":{"type":"none"}}`)}, nil + return sessions.Session{ID: command.SessionID, TenantID: command.TenantID, Metadata: command.Metadata, Configuration: json.RawMessage(`{"agent":{"id":"agent_validation","model":"validation-model"},"environment":{"type":"none"}}`)}, nil } func (s *validationStore) CreateSession(_ context.Context, tenant string, input sessions.CreateSession) (sessions.Session, error) { @@ -67,7 +67,7 @@ func (s *validationStore) serve(d *Dependencies, f *testFakes) { f.sessionAdmission.createSession = s.CreateSession f.agents.create, f.agents.update = s.CreateAgent, s.UpdateAgent f.vaults.createVault = s.CreateVault - f.sessions.getSession, f.sessions.updateSessionMetadata = nil, s.UpdateSessionMetadata + f.sessionsReader.getSession, f.sessions.updateSessionMetadata = nil, s.UpdateSessionMetadata f.environmentTemplates.create, f.environmentTemplates.update = s.CreateEnvironmentTemplate, s.UpdateEnvironmentTemplate } diff --git a/services/core/internal/deployment/errors.go b/services/core/internal/deployment/errors.go index 27252d080..1875e9716 100644 --- a/services/core/internal/deployment/errors.go +++ b/services/core/internal/deployment/errors.go @@ -30,6 +30,9 @@ var ( ErrNodeAddressMismatch = errors.New("sandbox node Core address differs from the public URL") // ErrNodeExists rejects an enrollment whose node ID is already in use. ErrNodeExists = errors.New("sandbox node ID is already enrolled") + // ErrSandboxResetSessionBusy reports a hosted Session an automatic reset + // does not archive yet because it has active work. + ErrSandboxResetSessionBusy = errors.New("the hosted Session is busy") ) // GenerationStaleError rejects a change whose expected generation is not the diff --git a/services/core/internal/deployment/fakes_test.go b/services/core/internal/deployment/fakes_test.go index 8113024ec..526945814 100644 --- a/services/core/internal/deployment/fakes_test.go +++ b/services/core/internal/deployment/fakes_test.go @@ -6,6 +6,7 @@ import ( "testing" "time" + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/adminaudit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/coremetrics" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" @@ -63,6 +64,7 @@ type fakeExecutionStorage struct { withReservation func(context.Context, AllocationKey, func(sessions.LockedSession, ReservationTx) error) error withAllocation func(context.Context, AllocationKey, func(AllocationTx) error) error withAllocationCleanup func(context.Context, AllocationKey, func(AllocationCleanupTx) error) error + withSessionArchive func(context.Context, string, string, func(context.Context, sessions.LockedSession, SessionArchiveTx) error) error clearWake func(context.Context, string, time.Time) error } @@ -537,6 +539,13 @@ func (f *fakeExecutionStorage) WithAllocationCleanup(ctx context.Context, key Al return f.withAllocationCleanup(ctx, key, apply) } +func (f *fakeExecutionStorage) WithSessionArchive(ctx context.Context, tenantID, sessionID string, apply func(context.Context, sessions.LockedSession, SessionArchiveTx) error) error { + if f.withSessionArchive == nil { + unexpected(f.t, "WithSessionArchive") + } + return f.withSessionArchive(ctx, tenantID, sessionID, apply) +} + func (f *fakeExecutionStorage) ClearWake(ctx context.Context, allocationID string, observed time.Time) error { if f.clearWake == nil { unexpected(f.t, "ClearWake") @@ -627,3 +636,60 @@ func (f *fakeReader) CountRetainedAllocations(ctx context.Context, installationI } return f.countRetainedAllocations(ctx, installationID) } + +// fakeSessionReader serves the Session reads the observation resolver makes; +// every other read fails the test. +type fakeSessionReader struct { + t testing.TB + getSession func(context.Context, string, string) (sessions.Session, error) + measuredSessionUsage func(context.Context, string, string) (json.RawMessage, error) +} + +func (f *fakeSessionReader) GetSession(ctx context.Context, tenantID, sessionID string) (sessions.Session, error) { + if f.getSession == nil { + unexpected(f.t, "GetSession") + } + return f.getSession(ctx, tenantID, sessionID) +} + +func (f *fakeSessionReader) MeasuredSessionUsage(ctx context.Context, tenantID, sessionID string) (json.RawMessage, error) { + if f.measuredSessionUsage == nil { + unexpected(f.t, "MeasuredSessionUsage") + } + return f.measuredSessionUsage(ctx, tenantID, sessionID) +} + +func (f *fakeSessionReader) ListSessions(context.Context, string, string, int, bool, *string) (sessions.Page, error) { + unexpected(f.t, "ListSessions") + return sessions.Page{}, nil +} + +func (f *fakeSessionReader) SessionStreamSnapshot(context.Context, string, string) (sessions.Session, int64, error) { + unexpected(f.t, "SessionStreamSnapshot") + return sessions.Session{}, 0, nil +} + +func (f *fakeSessionReader) SessionEventCursor(context.Context, string, string) (int64, error) { + unexpected(f.t, "SessionEventCursor") + return 0, nil +} + +func (f *fakeSessionReader) ListSessionEvents(context.Context, string, string, int64) ([]sessions.SessionChange, error) { + unexpected(f.t, "ListSessionEvents") + return nil, nil +} + +func (f *fakeSessionReader) GetTurnDiagnosticsSnapshot(context.Context, string, string, string) (sessions.TurnDiagnosticsSnapshot, error) { + unexpected(f.t, "GetTurnDiagnosticsSnapshot") + return sessions.TurnDiagnosticsSnapshot{}, nil +} + +func (f *fakeSessionReader) GetSessionExecutionConfiguration(context.Context, string, string) (v1.SessionExecutionConfiguration, error) { + unexpected(f.t, "GetSessionExecutionConfiguration") + return v1.SessionExecutionConfiguration{}, nil +} + +func (f *fakeSessionReader) GetManagedSessionArchive(context.Context, string, string) (sessions.ManagedArchive, error) { + unexpected(f.t, "GetManagedSessionArchive") + return sessions.ManagedArchive{}, nil +} diff --git a/services/core/internal/runtimeobs/storeresolver/resolver.go b/services/core/internal/deployment/observation.go similarity index 79% rename from services/core/internal/runtimeobs/storeresolver/resolver.go rename to services/core/internal/deployment/observation.go index 8f7e8e8dd..b4871624e 100644 --- a/services/core/internal/runtimeobs/storeresolver/resolver.go +++ b/services/core/internal/deployment/observation.go @@ -1,4 +1,4 @@ -package storeresolver +package deployment import ( "bytes" @@ -8,36 +8,26 @@ import ( "fmt" "io" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -type sessionStore interface { - GetSession(context.Context, string, string) (sessions.Session, error) - MeasuredSessionUsage(context.Context, string, string) (json.RawMessage, error) +// ObservationResolver resolves a Session to the Runtime observation target of +// its Environment and allocation, and lists the hosted Sessions to observe. +type ObservationResolver struct { + sessions sessions.SessionReader + reader Reader } -// allocationReader is the part of deployment.Reader the resolver reads. -type allocationReader interface { - EnvironmentAllocation(context.Context, deployment.AllocationKey) (deployment.Allocation, error) - ObservationSessions(context.Context, string, int) (deployment.ObservationSessionPage, error) -} - -type Resolver struct { - store sessionStore - allocations allocationReader -} - -func NewResolver(s sessionStore, allocations allocationReader) (*Resolver, error) { - if s == nil || allocations == nil { +func NewObservationResolver(sessionReader sessions.SessionReader, reader Reader) (*ObservationResolver, error) { + if sessionReader == nil || reader == nil { return nil, errors.New("Runtime observation store is required") } - return &Resolver{store: s, allocations: allocations}, nil + return &ObservationResolver{sessions: sessionReader, reader: reader}, nil } -func (r *Resolver) Resolve(ctx context.Context, tenantID, sessionID string) (runtimeobs.Target, error) { - session, err := r.store.GetSession(ctx, tenantID, sessionID) +func (r *ObservationResolver) Resolve(ctx context.Context, tenantID, sessionID string) (runtimeobs.Target, error) { + session, err := r.sessions.GetSession(ctx, tenantID, sessionID) if err != nil { return runtimeobs.Target{}, fmt.Errorf("resolve Runtime Session: %w", err) } @@ -52,7 +42,7 @@ func (r *Resolver) Resolve(ctx context.Context, tenantID, sessionID string) (run target := runtimeobs.Target{TenantID: session.TenantID, SessionID: session.ID, Mode: runtimeobs.Mode(configuration.Environment.Type)} // Telemetry counts measured usage continuously, including active Turns. // Public Session usage stays null until every root Turn ends measured. - measured, err := r.store.MeasuredSessionUsage(ctx, session.TenantID, session.ID) + measured, err := r.sessions.MeasuredSessionUsage(ctx, session.TenantID, session.ID) if err != nil { return runtimeobs.Target{}, fmt.Errorf("resolve Runtime Session usage: %w", err) } @@ -84,8 +74,8 @@ func (r *Resolver) Resolve(ctx context.Context, tenantID, sessionID string) (run return runtimeobs.Target{}, errors.New("managed Environment does not match resolved ownership") } target.EnvironmentID = session.Environment.ID - allocation, err := r.allocations.EnvironmentAllocation(ctx, deployment.AllocationKey{TenantID: tenantID, EnvironmentID: target.EnvironmentID}) - if errors.Is(err, deployment.ErrNotFound) { + allocation, err := r.reader.EnvironmentAllocation(ctx, AllocationKey{TenantID: tenantID, EnvironmentID: target.EnvironmentID}) + if errors.Is(err, ErrNotFound) { return target, runtimeobs.ErrUnavailable } if err != nil { @@ -153,8 +143,8 @@ func decodeTokenUsage(raw json.RawMessage) (*runtimeobs.TokenUsage, error) { return &runtimeobs.TokenUsage{InputTokens: uint64(*usage.InputTokens), OutputTokens: uint64(*usage.OutputTokens)}, nil } -func (r *Resolver) ListRuntimeObservationSessions(ctx context.Context, after string, limit int) (runtimeobs.SessionPage, error) { - page, err := r.allocations.ObservationSessions(ctx, after, limit) +func (r *ObservationResolver) ListRuntimeObservationSessions(ctx context.Context, after string, limit int) (runtimeobs.SessionPage, error) { + page, err := r.reader.ObservationSessions(ctx, after, limit) if err != nil { return runtimeobs.SessionPage{}, err } diff --git a/services/core/internal/runtimeobs/storeresolver/resolver_test.go b/services/core/internal/deployment/observation_test.go similarity index 70% rename from services/core/internal/runtimeobs/storeresolver/resolver_test.go rename to services/core/internal/deployment/observation_test.go index 593d6283a..b84bc1c57 100644 --- a/services/core/internal/runtimeobs/storeresolver/resolver_test.go +++ b/services/core/internal/deployment/observation_test.go @@ -1,4 +1,4 @@ -package storeresolver +package deployment import ( "context" @@ -6,44 +6,39 @@ import ( "errors" "testing" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) +// resolverStore is the Session and the deployment reads one resolution sees. type resolverStore struct { session sessions.Session measured json.RawMessage - allocation deployment.Allocation + allocation Allocation allocationErr error - page deployment.ObservationSessionPage + page ObservationSessionPage } -func (s resolverStore) GetSession(context.Context, string, string) (sessions.Session, error) { - return s.session, nil -} - -func (s resolverStore) MeasuredSessionUsage(_ context.Context, tenant, session string) (json.RawMessage, error) { - if tenant != s.session.TenantID || session != s.session.ID { - return nil, errors.New("measured usage read for another Session") +func newTestResolver(t *testing.T, s resolverStore) (*ObservationResolver, error) { + reader := &fakeReader{t: t, + environmentAllocation: func(context.Context, AllocationKey) (Allocation, error) { return s.allocation, s.allocationErr }, + observationSessions: func(context.Context, string, int) (ObservationSessionPage, error) { return s.page, nil }, } - return s.measured, nil -} - -func (s resolverStore) EnvironmentAllocation(context.Context, deployment.AllocationKey) (deployment.Allocation, error) { - return s.allocation, s.allocationErr -} - -func (s resolverStore) ObservationSessions(context.Context, string, int) (deployment.ObservationSessionPage, error) { - return s.page, nil + sessionReader := &fakeSessionReader{t: t, + getSession: func(context.Context, string, string) (sessions.Session, error) { return s.session, nil }, + measuredSessionUsage: func(_ context.Context, tenant, session string) (json.RawMessage, error) { + if tenant != s.session.TenantID || session != s.session.ID { + t.Fatalf("measured usage read for another Session: %s/%s", tenant, session) + } + return s.measured, nil + }, + } + return NewObservationResolver(sessionReader, reader) } -// newTestResolver serves the Session and the deployment reads from one fake. -func newTestResolver(s resolverStore) (*Resolver, error) { return NewResolver(s, s) } - -func TestResolverListsOnlyProviderNeutralSessionIdentity(t *testing.T) { - r, err := newTestResolver(resolverStore{page: deployment.ObservationSessionPage{ - Sessions: []deployment.ObservationSession{{TenantID: "tenant", SessionID: "session"}}, +func TestObservationResolverListsOnlyProviderNeutralSessionIdentity(t *testing.T) { + r, err := newTestResolver(t, resolverStore{page: ObservationSessionPage{ + Sessions: []ObservationSession{{TenantID: "tenant", SessionID: "session"}}, NextCursor: "session", }}) if err != nil { @@ -55,11 +50,11 @@ func TestResolverListsOnlyProviderNeutralSessionIdentity(t *testing.T) { } } -func TestResolverBindsManagedSessionEnvironmentAndAllocation(t *testing.T) { - r, err := newTestResolver(resolverStore{ +func TestObservationResolverBindsManagedSessionEnvironmentAndAllocation(t *testing.T) { + r, err := newTestResolver(t, resolverStore{ session: sessions.Session{ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"openai_hosted"}}`), Environment: &sessions.Environment{ID: "environment", TenantID: "tenant", SessionID: "session"}}, measured: []byte(`{"input_tokens":120,"input_tokens_details":{"cached_tokens":20},"output_tokens":30,"output_tokens_details":{"reasoning_tokens":10},"total_tokens":150}`), - allocation: deployment.Allocation{ + allocation: Allocation{ ID: "allocation", TenantID: "tenant", SessionID: "session", EnvironmentID: "environment", ProviderKey: "provider", DeviceID: "device", ComputePhase: "running", ComputeState: []byte(`{"current":{"name":"sandbox"}}`), }, @@ -85,7 +80,7 @@ func TestResolverBindsManagedSessionEnvironmentAndAllocation(t *testing.T) { } } -func TestResolverRejectsInvalidCanonicalSessionUsage(t *testing.T) { +func TestObservationResolverRejectsInvalidCanonicalSessionUsage(t *testing.T) { for _, usage := range []string{ `{"input_tokens":2,"input_tokens_details":{"cached_tokens":0},"output_tokens":3,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":4}`, `{}`, @@ -93,7 +88,7 @@ func TestResolverRejectsInvalidCanonicalSessionUsage(t *testing.T) { `{"input_tokens":0,"input_tokens_details":{"cached_tokens":-1},"output_tokens":0,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":0}`, `{"input_tokens":0,"input_tokens_details":{"cached_tokens":0},"output_tokens":0,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":0,"unknown":0}`, } { - resolver, err := newTestResolver(resolverStore{session: sessions.Session{ + resolver, err := newTestResolver(t, resolverStore{session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"none"}}`), }, measured: []byte(usage)}) if err != nil { @@ -105,8 +100,8 @@ func TestResolverRejectsInvalidCanonicalSessionUsage(t *testing.T) { } } -func TestResolverKeepsNullCanonicalSessionUsageAbsent(t *testing.T) { - resolver, err := newTestResolver(resolverStore{session: sessions.Session{ +func TestObservationResolverKeepsNullCanonicalSessionUsageAbsent(t *testing.T) { + resolver, err := newTestResolver(t, resolverStore{session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"none"}}`), }, measured: []byte(" \n null \t")}) if err != nil { @@ -120,7 +115,7 @@ func TestResolverKeepsNullCanonicalSessionUsageAbsent(t *testing.T) { // Telemetry reads measured usage, not the public Session value, which stays // null while a root Turn runs. -func TestResolverUsesMeasuredRatherThanPublicSessionUsage(t *testing.T) { +func TestObservationResolverUsesMeasuredRatherThanPublicSessionUsage(t *testing.T) { measured := `{"input_tokens":7,"input_tokens_details":{"cached_tokens":0},"output_tokens":3,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":10}` for _, test := range []struct { public, measured string @@ -129,7 +124,7 @@ func TestResolverUsesMeasuredRatherThanPublicSessionUsage(t *testing.T) { {"null", measured, &runtimeobs.TokenUsage{InputTokens: 7, OutputTokens: 3}}, {measured, "null", nil}, } { - resolver, err := newTestResolver(resolverStore{session: sessions.Session{ + resolver, err := newTestResolver(t, resolverStore{session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"none"}}`), Usage: []byte(test.public), }, measured: []byte(test.measured)}) if err != nil { @@ -142,7 +137,7 @@ func TestResolverUsesMeasuredRatherThanPublicSessionUsage(t *testing.T) { } } -func TestResolverKeepsUnsupportedModesDistinct(t *testing.T) { +func TestObservationResolverKeepsUnsupportedModesDistinct(t *testing.T) { for _, tc := range []struct { mode string environment *sessions.Environment @@ -150,7 +145,7 @@ func TestResolverKeepsUnsupportedModesDistinct(t *testing.T) { {mode: "none"}, {mode: "self_hosted", environment: &sessions.Environment{ID: "environment", TenantID: "tenant", SessionID: "session"}}, } { - r, err := newTestResolver(resolverStore{session: sessions.Session{ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"` + tc.mode + `"}}`), Environment: tc.environment}}) + r, err := newTestResolver(t, resolverStore{session: sessions.Session{ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"` + tc.mode + `"}}`), Environment: tc.environment}}) if err != nil { t.Fatal(err) } @@ -161,10 +156,10 @@ func TestResolverKeepsUnsupportedModesDistinct(t *testing.T) { } } -func TestResolverReportsManagedAllocationAsUnavailable(t *testing.T) { - r, err := newTestResolver(resolverStore{ +func TestObservationResolverReportsManagedAllocationAsUnavailable(t *testing.T) { + r, err := newTestResolver(t, resolverStore{ session: sessions.Session{ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"openai_hosted"}}`), Environment: &sessions.Environment{ID: "environment", TenantID: "tenant", SessionID: "session"}}, - allocationErr: deployment.ErrNotFound, + allocationErr: ErrNotFound, }) if err != nil { t.Fatal(err) @@ -175,13 +170,13 @@ func TestResolverReportsManagedAllocationAsUnavailable(t *testing.T) { } } -func TestResolverRejectsMismatchedEnvironmentOwnership(t *testing.T) { +func TestObservationResolverRejectsMismatchedEnvironmentOwnership(t *testing.T) { for _, mode := range []string{"self_hosted", "openai_hosted"} { for _, environment := range []sessions.Environment{ {ID: "environment", TenantID: "other", SessionID: "session"}, {ID: "environment", TenantID: "tenant", SessionID: "other"}, } { - resolver, err := newTestResolver(resolverStore{session: sessions.Session{ + resolver, err := newTestResolver(t, resolverStore{session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"` + mode + `"}}`), Environment: &environment, }}) @@ -195,19 +190,19 @@ func TestResolverRejectsMismatchedEnvironmentOwnership(t *testing.T) { } } -func TestResolverRejectsMismatchedAllocationOwnership(t *testing.T) { - base := deployment.Allocation{ +func TestObservationResolverRejectsMismatchedAllocationOwnership(t *testing.T) { + base := Allocation{ ID: "allocation", TenantID: "tenant", SessionID: "session", EnvironmentID: "environment", ProviderKey: "provider", DeviceID: "device", } - for _, mutate := range []func(*deployment.Allocation){ - func(value *deployment.Allocation) { value.TenantID = "other" }, - func(value *deployment.Allocation) { value.SessionID = "other" }, - func(value *deployment.Allocation) { value.EnvironmentID = "other" }, + for _, mutate := range []func(*Allocation){ + func(value *Allocation) { value.TenantID = "other" }, + func(value *Allocation) { value.SessionID = "other" }, + func(value *Allocation) { value.EnvironmentID = "other" }, } { allocation := base mutate(&allocation) - resolver, err := newTestResolver(resolverStore{ + resolver, err := newTestResolver(t, resolverStore{ session: sessions.Session{ ID: "session", TenantID: "tenant", Configuration: []byte(`{"environment":{"type":"openai_hosted"}}`), Environment: &sessions.Environment{ID: "environment", TenantID: "tenant", SessionID: "session"}, diff --git a/services/core/internal/deployment/session_archive.go b/services/core/internal/deployment/session_archive.go new file mode 100644 index 000000000..d2b2578fd --- /dev/null +++ b/services/core/internal/deployment/session_archive.go @@ -0,0 +1,144 @@ +package deployment + +import ( + "context" + "errors" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/adminaudit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// ArchiveSession ends the lifetime of a hosted Session's Environment while +// keeping the public Session, its history and persisted files. It cancels the +// Session's work, revokes the allocation's device and requests its cleanup; +// the lifecycle performs the provider work, and only its confirmation releases +// the allocation. expectedGeneration must be the current generation of the Web +// managed deployment. +func (e *ExecutionOperations) ArchiveSession(ctx context.Context, tenantID, sessionID string, expectedGeneration uint64) (sessions.ManagedArchive, error) { + return e.archiveSession(ctx, tenantID, sessionID, expectedGeneration, nil) +} + +// ArchiveResetSession archives a hosted Session for the running reset, which +// its generation and request time identify, so that a cancelled reset's +// candidates never affect its successor. An automatic reset skips a busy +// Session with ErrSandboxResetSessionBusy; an Environment that already ended +// only reports its disposal. The archive is audited with the reset's +// administrator source. +func (e *ExecutionOperations) ArchiveResetSession(ctx context.Context, tenantID, sessionID string, generation uint64, requestedAt time.Time) (sessions.ManagedArchive, error) { + return e.archiveSession(ctx, tenantID, sessionID, generation, &requestedAt) +} + +func (e *ExecutionOperations) archiveSession(ctx context.Context, tenantID, sessionID string, generation uint64, resetRequestedAt *time.Time) (sessions.ManagedArchive, error) { + var result sessions.ManagedArchive + err := e.storage.WithSessionArchive(ctx, tenantID, sessionID, func(ctx context.Context, locked sessions.LockedSession, tx SessionArchiveTx) error { + if err := locked.Public(); err != nil { + return err + } + d, err := tx.LoadDeployment() + if err != nil { + return err + } + if err := checkArchiveDeployment(d, generation); err != nil { + return err + } + environment, err := tx.LoadEnvironment(ctx) + if errors.Is(err, sessions.ErrNotFound) { + return sessions.ErrInvalidInput + } + if err != nil { + return err + } + if kind, err := sessions.EnvironmentType(environment.Configuration); err != nil || kind != "openai_hosted" { + return sessions.ErrInvalidInput + } + ended := environment.Status == "failed" || environment.Status == "expired" + if resetRequestedAt != nil { + if err := checkArchiveReset(d, *resetRequestedAt); err != nil { + return err + } + if d.Reset.Clear == ResetAuto { + busy, err := tx.LoadResetBusy() + if err != nil { + return err + } + if busy { + return ErrSandboxResetSessionBusy + } + } + if ended { + result, err = tx.LoadArchive(ctx) + return err + } + source, err := tx.LoadResetSource() + if err != nil { + return err + } + if source.ProjectID, err = tx.LoadProject(); err != nil { + return err + } + ctx = adminaudit.WithSource(ctx, source) + } + current, allocated, err := tx.FindAllocation(environment.ID) + if err != nil { + return err + } + live := allocated && current.State != "released" + if live && current.ProviderKey != d.InstallationID { + return ErrConflict + } + if err := sessions.TrackInputActivity(ctx, tx, func(ctx context.Context) error { + if !ended { + if err := tx.ExpireEnvironment(ctx, environment.ID); err != nil { + return err + } + } + return sessions.CancelWork(ctx, tx) + }); err != nil { + return err + } + if live { + if err := tx.RequestArchiveCleanup(current); err != nil { + return err + } + } else if !allocated { + if err := tx.ReleasePlacement(); err != nil { + return err + } + } + if err := tx.RecordArchiveAudit(ctx); err != nil { + return err + } + result, err = tx.LoadArchive(ctx) + return err + }) + if err != nil { + return sessions.ManagedArchive{}, err + } + return result, nil +} + +// checkArchiveDeployment rejects an archive against a deployment whose +// generation is not the expected one, that Web does not manage or that has no +// provider. +func checkArchiveDeployment(d Record, expectedGeneration uint64) error { + if d.Generation != expectedGeneration { + return &GenerationStaleError{CurrentGeneration: d.Generation} + } + if !d.WebManaged || d.InstallationID == "" { + return ErrConflict + } + if d.Provider == "" { + return ErrNotConfigured + } + return nil +} + +// checkArchiveReset rejects a reset archive unless the reset requested at +// requestedAt is still running. +func checkArchiveReset(d Record, requestedAt time.Time) error { + if d.Reset == nil || !d.Reset.RequestedAt.Equal(requestedAt) { + return ErrConflict + } + return nil +} diff --git a/services/core/internal/deployment/session_archive_test.go b/services/core/internal/deployment/session_archive_test.go new file mode 100644 index 000000000..3ebbd7cdb --- /dev/null +++ b/services/core/internal/deployment/session_archive_test.go @@ -0,0 +1,322 @@ +package deployment + +import ( + "context" + "encoding/json" + "errors" + "slices" + "strings" + "testing" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/adminaudit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +func TestCheckArchiveDeployment(t *testing.T) { + managed := Record{InstallationID: "installation", WebManaged: true, Provider: "docker", Generation: 4} + for name, test := range map[string]struct { + change func(*Record) + want error + }{ + "current": {func(*Record) {}, nil}, + "stale generation": {func(d *Record) { d.Generation = 5; d.WebManaged = false }, &GenerationStaleError{CurrentGeneration: 5}}, + "process managed": {func(d *Record) { d.WebManaged = false; d.Provider = "" }, ErrConflict}, + "no installation": {func(d *Record) { d.InstallationID = "" }, ErrConflict}, + "no provider": {func(d *Record) { d.Provider = "" }, ErrNotConfigured}, + } { + d := managed + test.change(&d) + err := checkArchiveDeployment(d, 4) + var stale *GenerationStaleError + switch want := test.want.(type) { + case nil: + if err != nil { + t.Errorf("%s: %v", name, err) + } + case *GenerationStaleError: + if !errors.As(err, &stale) || stale.CurrentGeneration != want.CurrentGeneration { + t.Errorf("%s: got %v, want stale generation %d", name, err, want.CurrentGeneration) + } + default: + if !errors.Is(err, want) || errors.As(err, &stale) { + t.Errorf("%s: got %v, want %v", name, err, want) + } + } + } +} + +func TestCheckArchiveReset(t *testing.T) { + requested := time.Unix(100, 0) + for name, test := range map[string]struct { + reset *ResetState + want error + }{ + "running reset": {&ResetState{Clear: ResetAuto, RequestedAt: requested}, nil}, + "no reset": {nil, ErrConflict}, + "another request": {&ResetState{Clear: ResetAuto, RequestedAt: requested.Add(time.Second)}, ErrConflict}, + } { + if err := checkArchiveReset(Record{Reset: test.reset}, requested); !errors.Is(err, test.want) { + t.Errorf("%s: got %v, want %v", name, err, test.want) + } + } +} + +// fakeArchiveTx is a Session archive over fixed facts. Each method records +// its call; a test asserts the exact call log. The Session's Turn and input +// reads fail the test unless it settles the Session with idle. +type fakeArchiveTx struct { + t testing.TB + calls []string + deployment Record + environment *sessions.Environment + busy bool + allocation *Allocation + loadActiveTurn func() (sessions.Turn, bool, error) + cancelPendingInput func() error + loadEnvironmentInput func() (*sessions.EnvironmentInputState, error) + // audited is the administrator source the audit was recorded with. + audited *adminaudit.Source +} + +// idle settles a Session with no active Turn, pending input or input +// activity change. +func (f *fakeArchiveTx) idle() { + f.loadActiveTurn = func() (sessions.Turn, bool, error) { return sessions.Turn{}, false, nil } + f.cancelPendingInput = func() error { return nil } + f.loadEnvironmentInput = func() (*sessions.EnvironmentInputState, error) { return nil, nil } +} + +func (f *fakeArchiveTx) record(call string) { f.calls = append(f.calls, call) } + +func (f *fakeArchiveTx) LoadDeployment() (Record, error) { + f.record("LoadDeployment") + return f.deployment, nil +} + +func (f *fakeArchiveTx) LoadEnvironment(context.Context) (sessions.Environment, error) { + f.record("LoadEnvironment") + if f.environment == nil { + return sessions.Environment{}, sessions.ErrNotFound + } + return *f.environment, nil +} + +func (f *fakeArchiveTx) ExpireEnvironment(_ context.Context, environment string) error { + f.record("ExpireEnvironment " + environment) + return nil +} + +func (f *fakeArchiveTx) LoadResetBusy() (bool, error) { + f.record("LoadResetBusy") + return f.busy, nil +} + +func (f *fakeArchiveTx) LoadResetSource() (adminaudit.Source, error) { + f.record("LoadResetSource") + return adminaudit.Source{CredentialID: "admin"}, nil +} + +func (f *fakeArchiveTx) LoadProject() (string, error) { + f.record("LoadProject") + return "project", nil +} + +func (f *fakeArchiveTx) FindAllocation(environment string) (Allocation, bool, error) { + f.record("FindAllocation " + environment) + if f.allocation == nil { + return Allocation{}, false, nil + } + return *f.allocation, true, nil +} + +func (f *fakeArchiveTx) RequestArchiveCleanup(current Allocation) error { + f.record("RequestArchiveCleanup " + current.DeviceID + " " + current.ID) + return nil +} + +func (f *fakeArchiveTx) ReleasePlacement() error { + f.record("ReleasePlacement") + return nil +} + +func (f *fakeArchiveTx) RecordArchiveAudit(ctx context.Context) error { + f.record("RecordArchiveAudit") + if source, ok := adminaudit.FromContext(ctx); ok { + f.audited = &source + } + return nil +} + +func (f *fakeArchiveTx) LoadArchive(context.Context) (sessions.ManagedArchive, error) { + f.record("LoadArchive") + return sessions.ManagedArchive{SessionID: "session", EnvironmentID: "environment", State: "cleanup_pending"}, nil +} + +func (f *fakeArchiveTx) LoadActiveTurn(context.Context) (sessions.Turn, bool, error) { + f.record("LoadActiveTurn") + if f.loadActiveTurn == nil { + unexpected(f.t, "LoadActiveTurn") + } + return f.loadActiveTurn() +} + +func (f *fakeArchiveTx) CancelPendingInput(context.Context) error { + f.record("CancelPendingInput") + if f.cancelPendingInput == nil { + unexpected(f.t, "CancelPendingInput") + } + return f.cancelPendingInput() +} + +func (f *fakeArchiveTx) LoadEnvironmentInput(context.Context) (*sessions.EnvironmentInputState, error) { + f.record("LoadEnvironmentInput") + if f.loadEnvironmentInput == nil { + unexpected(f.t, "LoadEnvironmentInput") + } + return f.loadEnvironmentInput() +} + +func (f *fakeArchiveTx) RequestTurnCancel(context.Context, string) error { + unexpected(f.t, "RequestTurnCancel") + return nil +} + +func (f *fakeArchiveTx) LoadTurn(context.Context, string) (sessions.Turn, error) { + unexpected(f.t, "LoadTurn") + return sessions.Turn{}, nil +} + +func (f *fakeArchiveTx) LoadEnding(context.Context, string) (sessions.Ending, error) { + unexpected(f.t, "LoadEnding") + return sessions.Ending{}, nil +} + +func (f *fakeArchiveTx) ApplyTurnEnd(context.Context, string, sessions.TurnEnd) error { + unexpected(f.t, "ApplyTurnEnd") + return nil +} + +func (f *fakeArchiveTx) LoadUsage(context.Context) (json.RawMessage, error) { + unexpected(f.t, "LoadUsage") + return nil, nil +} + +func (f *fakeArchiveTx) AppendChanges(context.Context, ...sessions.SessionChange) error { + unexpected(f.t, "AppendChanges") + return nil +} + +// archiveOperations builds execution operations whose Session archives run on +// tx with locked. +func archiveOperations(t *testing.T, tx *fakeArchiveTx, locked sessions.LockedSession) *ExecutionOperations { + t.Helper() + storage := &fakeExecutionStorage{t: t, withSessionArchive: func(ctx context.Context, tenantID, sessionID string, apply func(context.Context, sessions.LockedSession, SessionArchiveTx) error) error { + if tenantID != "tenant" || sessionID != "session" { + t.Fatalf("archived %s/%s", tenantID, sessionID) + } + return apply(ctx, locked, tx) + }} + operations, err := NewExecutionOperations(newService(t, &fakeStorage{t: t}, &fakeReader{t: t}, testPublicURL), storage) + if err != nil { + t.Fatal(err) + } + return operations +} + +func TestArchiveSession(t *testing.T) { + managed := Record{InstallationID: "installation", WebManaged: true, Provider: "docker", Generation: 1} + hosted := &sessions.Environment{ID: "environment", Status: "connected", Configuration: json.RawMessage(`{"type":"openai_hosted"}`)} + expired := &sessions.Environment{ID: "environment", Status: "expired", Configuration: hosted.Configuration} + live := &Allocation{ID: "allocation", DeviceID: "device", ProviderKey: "installation", State: "running"} + settle := []string{"LoadEnvironmentInput", "LoadActiveTurn", "CancelPendingInput", "LoadEnvironmentInput"} + expire := append([]string{"LoadEnvironmentInput", "ExpireEnvironment environment"}, settle[1:]...) + head := []string{"LoadDeployment", "LoadEnvironment", "FindAllocation environment"} + join := func(parts ...[]string) []string { + var calls []string + for _, part := range parts { + calls = append(calls, part...) + } + return calls + } + for _, test := range []struct { + name string + locked sessions.LockedSession + environment *sessions.Environment + allocation *Allocation + want error + calls []string + }{ + {"live allocation", sessions.LockedSession{}, hosted, live, nil, + join(head, expire, []string{"RequestArchiveCleanup device allocation", "RecordArchiveAudit", "LoadArchive"})}, + {"no allocation", sessions.LockedSession{}, hosted, nil, nil, join(head, expire, []string{"ReleasePlacement", "RecordArchiveAudit", "LoadArchive"})}, + {"released allocation", sessions.LockedSession{}, hosted, &Allocation{ID: "allocation", ProviderKey: "previous", State: "released"}, nil, + join(head, expire, []string{"RecordArchiveAudit", "LoadArchive"})}, + {"ended Environment", sessions.LockedSession{}, expired, live, nil, + join(head, settle, []string{"RequestArchiveCleanup device allocation", "RecordArchiveAudit", "LoadArchive"})}, + {"allocation of another installation", sessions.LockedSession{}, hosted, &Allocation{ID: "allocation", ProviderKey: "previous", State: "running"}, ErrConflict, head}, + {"deleted Session", sessions.LockedSession{Deleted: true}, hosted, nil, sessions.ErrNotFound, nil}, + {"no Environment", sessions.LockedSession{}, nil, nil, sessions.ErrInvalidInput, head[:2]}, + {"self-hosted Environment", sessions.LockedSession{}, &sessions.Environment{ID: "environment", Configuration: json.RawMessage(`{"type":"self_hosted"}`)}, nil, sessions.ErrInvalidInput, head[:2]}, + } { + t.Run(test.name, func(t *testing.T) { + tx := &fakeArchiveTx{t: t, deployment: managed, environment: test.environment, allocation: test.allocation} + if slices.Contains(test.calls, "LoadActiveTurn") { + tx.idle() + } + result, err := archiveOperations(t, tx, test.locked).ArchiveSession(t.Context(), "tenant", "session", 1) + if test.want == nil && (err != nil || result.State != "cleanup_pending") || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v %v, want %v", result, err, test.want) + } + if strings.Join(tx.calls, "\n") != strings.Join(test.calls, "\n") { + t.Fatalf("calls %q, want %q", tx.calls, test.calls) + } + if tx.audited != nil { + t.Fatal("an administrator archive borrowed a reset source", tx.audited) + } + }) + } +} + +func TestArchiveResetSession(t *testing.T) { + requested := time.Unix(100, 0) + resetting := func(clear string) Record { + return Record{InstallationID: "installation", WebManaged: true, Provider: "docker", Generation: 1, Reset: &ResetState{Clear: clear, RequestedAt: requested}} + } + hosted := &sessions.Environment{ID: "environment", Status: "connected", Configuration: json.RawMessage(`{"type":"openai_hosted"}`)} + failed := &sessions.Environment{ID: "environment", Status: "failed", Configuration: hosted.Configuration} + for _, test := range []struct { + name string + deployment Record + environment *sessions.Environment + busy bool + requested time.Time + want error + calls []string + }{ + {"idle Session", resetting(ResetAuto), hosted, false, requested, nil, + []string{"LoadDeployment", "LoadEnvironment", "LoadResetBusy", "LoadResetSource", "LoadProject", "FindAllocation environment", "LoadEnvironmentInput", "ExpireEnvironment environment", "LoadActiveTurn", "CancelPendingInput", "LoadEnvironmentInput", "ReleasePlacement", "RecordArchiveAudit", "LoadArchive"}}, + {"busy Session", resetting(ResetAuto), hosted, true, requested, ErrSandboxResetSessionBusy, []string{"LoadDeployment", "LoadEnvironment", "LoadResetBusy"}}, + {"forced busy Session", resetting(ResetForce), hosted, true, requested, nil, + []string{"LoadDeployment", "LoadEnvironment", "LoadResetSource", "LoadProject", "FindAllocation environment", "LoadEnvironmentInput", "ExpireEnvironment environment", "LoadActiveTurn", "CancelPendingInput", "LoadEnvironmentInput", "ReleasePlacement", "RecordArchiveAudit", "LoadArchive"}}, + {"ended Environment", resetting(ResetAuto), failed, false, requested, nil, []string{"LoadDeployment", "LoadEnvironment", "LoadResetBusy", "LoadArchive"}}, + {"another reset", resetting(ResetAuto), hosted, false, requested.Add(time.Second), ErrConflict, []string{"LoadDeployment", "LoadEnvironment"}}, + } { + t.Run(test.name, func(t *testing.T) { + tx := &fakeArchiveTx{t: t, deployment: test.deployment, environment: test.environment, busy: test.busy} + if slices.Contains(test.calls, "LoadActiveTurn") { + tx.idle() + } + _, err := archiveOperations(t, tx, sessions.LockedSession{}).ArchiveResetSession(t.Context(), "tenant", "session", 1, test.requested) + if test.want == nil && err != nil || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v, want %v", err, test.want) + } + if strings.Join(tx.calls, "\n") != strings.Join(test.calls, "\n") { + t.Fatalf("calls %q, want %q", tx.calls, test.calls) + } + if audited := tx.audited; test.want == nil && strings.Contains(strings.Join(test.calls, ","), "RecordArchiveAudit") && (audited == nil || audited.CredentialID != "admin" || audited.ProjectID != "project") { + t.Fatal("reset archive was not audited with the reset source", audited) + } + }) + } +} diff --git a/services/core/internal/deployment/storage.go b/services/core/internal/deployment/storage.go index 399c98d88..20cff0768 100644 --- a/services/core/internal/deployment/storage.go +++ b/services/core/internal/deployment/storage.go @@ -61,6 +61,11 @@ type ExecutionStorage interface { // missing allocation is ErrNotFound. WithAllocation(ctx context.Context, key AllocationKey, apply func(AllocationTx) error) error AllocationCleanupStorage + // WithSessionArchive runs apply in one leased transaction that locks the + // tenant's Session, deleted or not, and then prunes the Session's journal. + // It commits only when both succeed. A malformed tenant is + // ErrInvalidInput; a malformed or missing Session is sessions.ErrNotFound. + WithSessionArchive(ctx context.Context, tenantID, sessionID string, apply func(context.Context, sessions.LockedSession, SessionArchiveTx) error) error // ClearWake clears the allocation's wake request unless activity newer // than observed arrived. It runs on the lease. ClearWake(ctx context.Context, allocationID string, observed time.Time) error @@ -133,6 +138,47 @@ type AllocationCleanupTx interface { RequestCleanup(current Allocation) (Allocation, error) } +// SessionArchiveTx is one Session-locked archive of the Session's hosted +// Environment, with the Session's cancellation bound to the same transaction. +type SessionArchiveTx interface { + sessions.CancellationTx + sessions.InputActivityTx + // LoadDeployment locks the deployment and returns it. The Session lock + // precedes it, as in Turn, allocation and input admission. + LoadDeployment() (Record, error) + // LoadEnvironment reads the Session's Environment. + LoadEnvironment(ctx context.Context) (sessions.Environment, error) + // ExpireEnvironment sets the status of the Session's Environment to + // expired. + ExpireEnvironment(ctx context.Context, environment string) error + // LoadResetBusy reports whether the Session has an in-progress or waiting + // Turn or subagent Turn, or a pending Environment file write, which an + // automatic reset does not archive. + LoadResetBusy() (bool, error) + // LoadResetSource returns the administrator source that started the + // running reset. + LoadResetSource() (adminaudit.Source, error) + // LoadProject returns the ID of the tenant's Project. + LoadProject() (string, error) + // FindAllocation returns the Environment's allocation and whether it has + // one. + FindAllocation(environment string) (Allocation, bool, error) + // RequestArchiveCleanup revokes the allocation's device and records that + // its resources await cleanup. The device's first revocation records the + // Session's active Turn whose cancellation was requested, which the + // archived cancellation receipt reports. + RequestArchiveCleanup(current Allocation) error + // ReleasePlacement releases the node placement of the Session's + // Environments that have no allocation. + ReleasePlacement() error + // RecordArchiveAudit records the administrator audit of the archive with + // the source on ctx. + RecordArchiveAudit(ctx context.Context) error + // LoadArchive reads the resource disposal of the Session's hosted + // Environment. + LoadArchive(ctx context.Context) (sessions.ManagedArchive, error) +} + // Reader answers deployment and node queries. type Reader interface { // Deployment returns the stored deployment. diff --git a/services/core/internal/execution/archive_cancellation_cleanup_test.go b/services/core/internal/execution/archive_cancellation_cleanup_test.go index eda95d545..822676a52 100644 --- a/services/core/internal/execution/archive_cancellation_cleanup_test.go +++ b/services/core/internal/execution/archive_cancellation_cleanup_test.go @@ -148,7 +148,7 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { } defer release() } - if _, err := writer.ArchiveManagedSession(audit, project.TenantID, session.ID, 1); err != nil { + if _, err := leased.Deployment.ArchiveSession(audit, project.TenantID, session.ID, 1); err != nil { t.Fatal(err) } owner, err = reader.EnvironmentAllocation(t.Context(), key) diff --git a/services/core/internal/execution/deployment_fixture_test.go b/services/core/internal/execution/deployment_fixture_test.go index 2c21c8010..94bacf49e 100644 --- a/services/core/internal/execution/deployment_fixture_test.go +++ b/services/core/internal/execution/deployment_fixture_test.go @@ -255,6 +255,10 @@ func (s *strictExecutionStorage) WithAllocationCleanup(ctx context.Context, key return s.withAllocationCleanup(ctx, key, apply) } +func (s *strictExecutionStorage) WithSessionArchive(context.Context, string, string, func(context.Context, sessions.LockedSession, deployment.SessionArchiveTx) error) error { + return unexpectedDeploymentCall(s.t, "WithSessionArchive") +} + func (s *strictExecutionStorage) ClearWake(ctx context.Context, allocationID string, observed time.Time) error { if s.clearWake == nil { return unexpectedDeploymentCall(s.t, "ClearWake") diff --git a/services/core/internal/execution/dispatcher.go b/services/core/internal/execution/dispatcher.go index 12710583e..d1e06b423 100644 --- a/services/core/internal/execution/dispatcher.go +++ b/services/core/internal/execution/dispatcher.go @@ -75,7 +75,7 @@ type Result struct { // Run claims once before subscribing or sending. Uncertain deliveries are not replayed. func (d *Dispatcher) Run(ctx context.Context, tenantID, sessionID, turnID string) (sessions.Turn, error) { - session, err := d.Store.GetSession(ctx, tenantID, sessionID) + session, err := d.SessionsReader.GetSession(ctx, tenantID, sessionID) if err != nil { return sessions.Turn{}, err } diff --git a/services/core/internal/execution/environment_directory.go b/services/core/internal/execution/environment_directory.go index bba21b5dd..5d1286efb 100644 --- a/services/core/internal/execution/environment_directory.go +++ b/services/core/internal/execution/environment_directory.go @@ -85,7 +85,7 @@ func (w *Worker) runDirectoryRead(owner context.Context, request directoryReadRe if err != nil { return } - session, err := w.dispatcher.Store.GetSession(check, environment.TenantID, environment.SessionID) + session, err := w.dispatcher.SessionsReader.GetSession(check, environment.TenantID, environment.SessionID) if err != nil { result.err = err return diff --git a/services/core/internal/execution/environment_file_write.go b/services/core/internal/execution/environment_file_write.go index 57d06f99e..d608265fd 100644 --- a/services/core/internal/execution/environment_file_write.go +++ b/services/core/internal/execution/environment_file_write.go @@ -85,7 +85,7 @@ func (w *Worker) runFileWrite(owner context.Context, request fileWriteRequest) f if err != nil || (placement.Type != "openai_hosted" && placement.Type != "self_hosted") { return unavailable } - session, err := w.dispatcher.Store.GetSession(ctx, environment.TenantID, environment.SessionID) + session, err := w.dispatcher.SessionsReader.GetSession(ctx, environment.TenantID, environment.SessionID) if err != nil { return fileWriteResult{err: err} } diff --git a/services/core/internal/execution/prepared_dispatch.go b/services/core/internal/execution/prepared_dispatch.go index c9dec7ede..1584b7442 100644 --- a/services/core/internal/execution/prepared_dispatch.go +++ b/services/core/internal/execution/prepared_dispatch.go @@ -26,7 +26,7 @@ func (d *Dispatcher) RunEnvironmentInput(ctx context.Context, lease Ownership, t if err != nil || run.Reservation.State != sessions.EnvironmentInputPending { return run, err } - session, err := d.Store.GetSession(ctx, tenantID, sessionID) + session, err := d.SessionsReader.GetSession(ctx, tenantID, sessionID) if err != nil { return run, err } diff --git a/services/core/internal/execution/sandbox_reset.go b/services/core/internal/execution/sandbox_reset.go index 187138fd4..f841d4ee8 100644 --- a/services/core/internal/execution/sandbox_reset.go +++ b/services/core/internal/execution/sandbox_reset.go @@ -8,7 +8,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/obs/log" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func (w *Worker) StartSandboxReset(ctx context.Context, input deployment.ResetRequest) (deployment.View, error) { @@ -106,9 +105,9 @@ func (m *runtimeManager) resetPage(parent, ctx context.Context) error { if ctx.Err() != nil { return nil } - _, err := m.store.ArchiveSandboxResetSession(ctx, candidate.TenantID, candidate.SessionID, current.Generation, current.Reset.RequestedAt) + _, err := m.deployment.ArchiveResetSession(ctx, candidate.TenantID, candidate.SessionID, current.Generation, current.Reset.RequestedAt) m.resetCursor = candidate.SessionID - if err != nil && !errors.Is(err, store.ErrSandboxResetSessionBusy) && !errors.Is(err, sessions.ErrNotFound) { + if err != nil && !errors.Is(err, deployment.ErrSandboxResetSessionBusy) && !errors.Is(err, sessions.ErrNotFound) { // Do not log a provider body, request, credential or stored provenance. log.Warn(ctx, "Sandbox reset archive remains pending", "session_id", candidate.SessionID) } diff --git a/services/core/internal/execution/session_archive.go b/services/core/internal/execution/session_archive.go deleted file mode 100644 index 65ed5ba16..000000000 --- a/services/core/internal/execution/session_archive.go +++ /dev/null @@ -1,13 +0,0 @@ -package execution - -import ( - "context" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -// ArchiveManagedSession records administrative cleanup intent through the same -// execution owner that settles resources. Provider work remains in the lifecycle. -func (w *Worker) ArchiveManagedSession(ctx context.Context, tenant, session string, generation uint64) (sessions.ManagedArchive, error) { - return w.dispatcher.Store.ArchiveManagedSession(ctx, tenant, session, generation) -} diff --git a/services/core/internal/execution/worker.go b/services/core/internal/execution/worker.go index 887722017..0b2b48a0b 100644 --- a/services/core/internal/execution/worker.go +++ b/services/core/internal/execution/worker.go @@ -122,7 +122,7 @@ func (w *Worker) CheckOwnership(ctx context.Context) error { } func (w *Worker) SubmitInputs(ctx context.Context, tenant, session, key string, inputs []sessions.Input) ([]sessions.InputReceipt, error) { - value, err := w.admission.GetSession(ctx, tenant, session) + value, err := w.dispatcher.SessionsReader.GetSession(ctx, tenant, session) if err != nil { return nil, err } diff --git a/services/core/internal/execution/worker_device.go b/services/core/internal/execution/worker_device.go index ea63a6b9f..015ac962e 100644 --- a/services/core/internal/execution/worker_device.go +++ b/services/core/internal/execution/worker_device.go @@ -46,7 +46,7 @@ func (w *Worker) bind(ctx context.Context, item sessions.ExecutionWork) (bool, e } func (w *Worker) bindDevice(ctx context.Context, tenantID, sessionID string, input proto.MessageInput) (bool, error) { - session, err := w.dispatcher.Store.GetSession(ctx, tenantID, sessionID) + session, err := w.dispatcher.SessionsReader.GetSession(ctx, tenantID, sessionID) if errors.Is(err, sessions.ErrNotFound) { return false, nil } diff --git a/services/core/internal/persistence/postgres/deploymentpg/allocations.go b/services/core/internal/persistence/postgres/deploymentpg/allocations.go index c12d7bca5..5ff80fbec 100644 --- a/services/core/internal/persistence/postgres/deploymentpg/allocations.go +++ b/services/core/internal/persistence/postgres/deploymentpg/allocations.go @@ -30,6 +30,19 @@ func allocation(row sqlc.RuntimeAllocation, session, tenant pgtype.UUID, deleted } // allocationKey returns the stored form of the key's identifiers. +// findAllocation reads the tenant's allocation of the Environment; found is +// false when it has none. +func findAllocation(ctx context.Context, q *sqlc.Queries, tenant, environment pgtype.UUID) (deployment.Allocation, bool, error) { + row, err := q.GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: tenant, EnvironmentID: environment}) + if errors.Is(err, pgx.ErrNoRows) { + return deployment.Allocation{}, false, nil + } + if err != nil { + return deployment.Allocation{}, false, err + } + return allocation(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired), true, nil +} + func allocationKey(key deployment.AllocationKey) (pgtype.UUID, pgtype.UUID, error) { tenant, err := parseID(key.TenantID) if err != nil { @@ -175,14 +188,7 @@ type reservationTx struct { } func (t *reservationTx) FindAllocation() (deployment.Allocation, bool, error) { - row, err := t.q.GetRuntimeAllocation(t.ctx, sqlc.GetRuntimeAllocationParams{TenantID: t.tenant, EnvironmentID: t.environment}) - if errors.Is(err, pgx.ErrNoRows) { - return deployment.Allocation{}, false, nil - } - if err != nil { - return deployment.Allocation{}, false, err - } - return allocation(row.RuntimeAllocation, t.session, t.tenant, row.DeletedAt, row.Expired), true, nil + return findAllocation(t.ctx, t.q, t.tenant, t.environment) } func (t *reservationTx) LockDeployment() (placement.Deployment, error) { @@ -229,14 +235,11 @@ type allocationTx struct { } func (t *allocationTx) LoadAllocation() (deployment.Allocation, error) { - row, err := t.q.GetRuntimeAllocation(t.ctx, sqlc.GetRuntimeAllocationParams{TenantID: t.tenant, EnvironmentID: t.environment}) - if errors.Is(err, pgx.ErrNoRows) { - return deployment.Allocation{}, deployment.ErrNotFound + current, found, err := findAllocation(t.ctx, t.q, t.tenant, t.environment) + if err == nil && !found { + err = deployment.ErrNotFound } - if err != nil { - return deployment.Allocation{}, err - } - return allocation(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired), nil + return current, err } func (t *allocationTx) LoadSessionDevice() (deployment.SessionDevice, bool, error) { @@ -363,14 +366,11 @@ func (s *Store) EnvironmentAllocation(ctx context.Context, key deployment.Alloca if err != nil { return deployment.Allocation{}, err } - row, err := s.pool.Queries().GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: tenant, EnvironmentID: environment}) - if errors.Is(err, pgx.ErrNoRows) { - return deployment.Allocation{}, deployment.ErrNotFound - } - if err != nil { - return deployment.Allocation{}, err + current, found, err := findAllocation(ctx, s.pool.Queries(), tenant, environment) + if err == nil && !found { + err = deployment.ErrNotFound } - return allocation(row.RuntimeAllocation, row.SessionID, row.TenantID, row.DeletedAt, row.Expired), nil + return current, err } func (s *Store) CredentialAllocations(ctx context.Context, after string) ([]deployment.Allocation, error) { diff --git a/services/core/internal/persistence/postgres/deploymentpg/session_archive.go b/services/core/internal/persistence/postgres/deploymentpg/session_archive.go new file mode 100644 index 000000000..a361597ee --- /dev/null +++ b/services/core/internal/persistence/postgres/deploymentpg/session_archive.go @@ -0,0 +1,85 @@ +package deploymentpg + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/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/sessions" +) + +func (e *Execution) WithSessionArchive(ctx context.Context, tenantID, sessionID string, apply func(context.Context, sessions.LockedSession, deployment.SessionArchiveTx) error) error { + tenant, err := parseID(tenantID) + if err != nil { + return err + } + session := pgunit.PathID(sessionID) + return e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + locked, err := sessionpg.LockSession(ctx, q, tenant, session) + if err != nil { + return err + } + t := &archiveTx{SessionTx: sessionpg.BindSession(q, tenant, session), unit: unit{ctx: ctx, q: q, locked: true}, tenantID: tenantID, tenant: tenant, session: session} + if err := apply(ctx, locked, t); err != nil { + return err + } + return sessionpg.PruneChanges(ctx, q, session) + }) +} + +// archiveTx is one Session archive. Its unit locks the deployment when it +// reads it; tenantID is the caller's tenant as the audit records it. +type archiveTx struct { + *sessionpg.SessionTx + unit + tenantID string + tenant, session pgtype.UUID +} + +func (t *archiveTx) LoadResetBusy() (bool, error) { + return t.q.SessionBlocksAutoReset(t.ctx, t.session) +} + +func (t *archiveTx) LoadProject() (string, error) { + project, err := t.q.GetSandboxResetProject(t.ctx, t.tenant) + return uuidString(project), err +} + +func (t *archiveTx) FindAllocation(environment string) (deployment.Allocation, bool, error) { + id, err := parseID(environment) + if err != nil { + return deployment.Allocation{}, false, err + } + return findAllocation(t.ctx, t.q, t.tenant, id) +} + +func (t *archiveTx) RequestArchiveCleanup(current deployment.Allocation) error { + device, err := parseID(current.DeviceID) + if err != nil { + return err + } + id, err := parseID(current.ID) + if err != nil { + return err + } + if _, err := t.q.RevokeArchivedRuntimeDevice(t.ctx, sqlc.RevokeArchivedRuntimeDeviceParams{TenantID: t.tenant, DeviceID: device, SessionID: t.session}); err != nil { + return err + } + _, err = t.q.RequestRuntimeCleanup(t.ctx, id) + return err +} + +func (t *archiveTx) ReleasePlacement() error { + return t.q.ReleaseUnallocatedRuntimePlacement(t.ctx, t.session) +} + +func (t *archiveTx) RecordArchiveAudit(ctx context.Context) error { + return auditpg.RecordAdminMutation(ctx, t.q, t.tenantID, "archive", "session", uuidString(t.session)) +} diff --git a/services/core/internal/persistence/postgres/deploymentpg/tx.go b/services/core/internal/persistence/postgres/deploymentpg/tx.go index fbab74148..6145dd13b 100644 --- a/services/core/internal/persistence/postgres/deploymentpg/tx.go +++ b/services/core/internal/persistence/postgres/deploymentpg/tx.go @@ -45,6 +45,23 @@ func (u unit) loadDeployment() (sqlc.RuntimeDeployment, error) { return u.q.GetRuntimeDeployment(u.ctx) } +var errNoResetSource = errors.New("sandbox reset has no stored administrator source") + +// LoadResetSource decodes the administrator source that started the running +// reset from the stored deployment. Audit entries the reset records later +// carry that source. +func (u unit) LoadResetSource() (adminaudit.Source, error) { + d, err := u.loadDeployment() + if err != nil { + return adminaudit.Source{}, err + } + var source adminaudit.Source + if !d.ResetClear.Valid || json.Unmarshal(d.ResetAudit, &source) != nil { + return adminaudit.Source{}, errNoResetSource + } + return source, nil +} + func (u unit) LoadNode(id string) (deployment.StoredNode, error) { nodeID, err := parseID(id) if err != nil { @@ -339,27 +356,6 @@ func (t *deploymentTx) ForceReset() error { return t.q.ForceSandboxReset(t.ctx) func (t *deploymentTx) CancelReset() error { return t.q.CancelSandboxReset(t.ctx) } -func (t *deploymentTx) LoadResetSource() (adminaudit.Source, error) { - d, err := t.loadDeployment() - if err != nil { - return adminaudit.Source{}, err - } - return ResetSource(d) -} - -var errNoResetSource = errors.New("sandbox reset has no stored administrator source") - -// ResetSource decodes the administrator source that started the running reset -// from the stored deployment row. Audit entries the reset records later carry -// that source. -func ResetSource(d sqlc.RuntimeDeployment) (adminaudit.Source, error) { - var source adminaudit.Source - if !d.ResetClear.Valid || json.Unmarshal(d.ResetAudit, &source) != nil { - return adminaudit.Source{}, errNoResetSource - } - return source, nil -} - func (t *deploymentTx) CompleteReset() error { if err := t.q.CompleteSandboxReset(t.ctx); err != nil { return err diff --git a/services/core/internal/persistence/postgres/sessionpg/binding.go b/services/core/internal/persistence/postgres/sessionpg/binding.go index cf227e3f5..745e40c12 100644 --- a/services/core/internal/persistence/postgres/sessionpg/binding.go +++ b/services/core/internal/persistence/postgres/sessionpg/binding.go @@ -196,3 +196,9 @@ func (t *SessionTx) LoadComputeSuspension(ctx context.Context) (bool, error) { func (t *SessionTx) LoadPendingFileWrite(ctx context.Context) (bool, error) { return t.q.EnvironmentFileWriteBlocksSession(ctx, t.session) } + +// LoadArchive reads the resource disposal of the Session's hosted +// Environment, as GetManagedSessionArchive does. +func (t *SessionTx) LoadArchive(ctx context.Context) (sessions.ManagedArchive, error) { + return loadManagedArchive(ctx, t.q, t.tenant, t.session) +} diff --git a/services/core/internal/persistence/postgres/sessionpg/session_reads.go b/services/core/internal/persistence/postgres/sessionpg/session_reads.go new file mode 100644 index 000000000..a32a6f622 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/session_reads.go @@ -0,0 +1,325 @@ +package sessionpg + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var _ sessions.SessionReader = (*Store)(nil) + +func (s *Store) GetSession(ctx context.Context, tenantID, sessionID string) (sessions.Session, error) { + session, _, err := s.SessionStreamSnapshot(ctx, tenantID, sessionID) + return session, err +} + +func (s *Store) ListSessions(ctx context.Context, tenantID, cursor string, limit int, ascending bool, agentID *string) (sessions.Page, error) { + tenant, err := parseID(tenantID) + if err != nil { + return sessions.Page{}, err + } + if limit < 1 || limit > 100 { + return sessions.Page{}, fmt.Errorf("%w: page size must be 1..100", sessions.ErrInvalidInput) + } + params := sqlc.ListSessionsParams{TenantID: tenant, PageLimit: int32(limit + 1), AfterID: pgtype.UUID{Valid: true}, Ascending: ascending} + if agentID != nil { + params.AgentID = pgtype.Text{String: *agentID, Valid: true} + } + var page sessions.Page + err = s.units.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + if cursor != "" { + after, err := q.GetSession(ctx, sqlc.GetSessionParams{TenantID: tenant, ID: pgunit.PathID(cursor)}) + if err != nil { + return err + } + params.AfterCreated, params.AfterID = after.CreatedAt, after.ID + } + rows, err := q.ListSessions(ctx, params) + if err != nil { + return fmt.Errorf("list sessions: %w", err) + } + page.Sessions = make([]sessions.Session, 0, min(limit, len(rows))) + if len(rows) > limit { + page.NextCursor = uuid.UUID(rows[limit-1].ID.Bytes).String() + rows = rows[:limit] + } + for _, row := range rows { + session, err := SessionFromRow(row) + if err == nil { + session, err = LoadSessionActivity(ctx, q, session) + } + if err != nil { + return err + } + page.Sessions = append(page.Sessions, session) + } + return nil + }) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.Page{}, sessions.ErrNotFound + } + return page, err +} + +func (s *Store) SessionStreamSnapshot(ctx context.Context, tenantID, sessionID string) (sessions.Session, int64, error) { + tenant, err := parseID(tenantID) + if err != nil { + return sessions.Session{}, 0, err + } + var session sessions.Session + var cursor int64 + err = s.units.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + row, err := q.GetSession(ctx, sqlc.GetSessionParams{TenantID: tenant, ID: pgunit.PathID(sessionID)}) + if err != nil { + return err + } + if session, err = SessionFromRow(row); err != nil { + return err + } + cursor = row.EventSequence + session, err = LoadSessionActivity(ctx, q, session) + return err + }) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.Session{}, 0, sessions.ErrNotFound + } + if err != nil { + return sessions.Session{}, 0, fmt.Errorf("get session: %w", err) + } + return session, cursor, nil +} + +func (s *Store) SessionEventCursor(ctx context.Context, tenantID, sessionID string) (int64, error) { + tenant, err := parseID(tenantID) + if err != nil { + return 0, err + } + cursor, err := s.units.Queries().SessionEventCursor(ctx, sqlc.SessionEventCursorParams{TenantID: tenant, ID: pgunit.PathID(sessionID)}) + if errors.Is(err, pgx.ErrNoRows) { + return 0, sessions.ErrNotFound + } + return cursor, err +} + +func (s *Store) ListSessionEvents(ctx context.Context, tenantID, sessionID string, after int64) ([]sessions.SessionChange, error) { + if after < 0 { + return nil, sessions.ErrInvalidInput + } + latest, err := s.SessionEventCursor(ctx, tenantID, sessionID) + if err != nil { + return nil, err + } + tenant, err := parseID(tenantID) + if err != nil { + return nil, err + } + id, err := parseID(sessionID) + if err != nil { + return nil, err + } + rows, err := s.units.Queries().ListSessionEvents(ctx, sqlc.ListSessionEventsParams{TenantID: tenant, SessionID: id, Sequence: after}) + if err != nil { + return nil, err + } + changes := make([]sessions.SessionChange, 0, len(rows)) + if len(rows) == 0 && latest > after { + return nil, sessions.ErrStreamGap + } + for _, row := range rows { + if row.Sequence != after+1 { + return nil, sessions.ErrStreamGap + } + var change sessions.SessionChange + decoder := json.NewDecoder(bytes.NewReader(row.Payload)) + decoder.UseNumber() + if err := decoder.Decode(&change); err != nil { + return nil, err + } + change.Sequence = row.Sequence + changes = append(changes, change) + after = row.Sequence + } + return changes, nil +} + +func (s *Store) GetTurnDiagnosticsSnapshot(ctx context.Context, tenantID, sessionID, turnID string) (sessions.TurnDiagnosticsSnapshot, error) { + params, err := publicTurnLookup(tenantID, sessionID, turnID) + if err != nil { + return sessions.TurnDiagnosticsSnapshot{}, err + } + result := sessions.TurnDiagnosticsSnapshot{Items: []sessions.ItemDiagnosticTiming{}} + err = s.units.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + turn, err := q.GetTurn(ctx, params) + if err != nil { + return err + } + row, err := q.GetSession(ctx, sqlc.GetSessionParams{TenantID: params.TenantID, ID: params.SessionID}) + if err != nil { + return err + } + result.Session, err = SessionFromRow(row) + if err != nil { + return err + } + result.Turn = TurnFromRow(turn) + rows, err := q.ListTurnItemDiagnostics(ctx, sqlc.ListTurnItemDiagnosticsParams{SessionID: params.SessionID, TurnID: params.ID}) + if err != nil { + return err + } + result.ItemsTruncated = len(rows) > 1000 + if result.ItemsTruncated { + rows = rows[:1000] + } + for _, row := range rows { + item := sessions.ItemDiagnosticTiming{ItemID: uuid.UUID(row.ID.Bytes).String(), StartedAt: row.CreatedAt.Time} + if row.SettledAt.Valid { + at := row.SettledAt.Time + item.CompletedAt = &at + } + result.Items = append(result.Items, item) + } + return nil + }) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.TurnDiagnosticsSnapshot{}, sessions.ErrNotFound + } + return result, err +} + +func (s *Store) GetSessionExecutionConfiguration(ctx context.Context, tenantID, sessionID string) (v1.SessionExecutionConfiguration, error) { + tenant, err := parseID(tenantID) + if err != nil { + return v1.SessionExecutionConfiguration{}, err + } + row, err := s.units.Queries().GetSessionExecutionConfiguration(ctx, sqlc.GetSessionExecutionConfigurationParams{TenantID: tenant, SessionID: pgunit.PathID(sessionID)}) + if errors.Is(err, pgx.ErrNoRows) { + return v1.SessionExecutionConfiguration{}, sessions.ErrNotFound + } + if err != nil { + return v1.SessionExecutionConfiguration{}, fmt.Errorf("get session execution configuration: %w", err) + } + var projection v1.SessionExecutionConfiguration + if len(row.ExecutionConfiguration) == 0 { + model, err := sessions.ExecutionModel(row.SessionConfiguration) + if err != nil { + return projection, err + } + var harness *string + if row.Engine != "" { + harness = &row.Engine + } + projection.Model = v1.ExecutionSelection{Value: model, Source: "unknown"} + projection.Harness = v1.ExecutionSelection{Value: harness, Source: "unknown"} + projection.ModelProvider = v1.ExecutionProviderSelection{Source: "unknown", Status: "unavailable"} + } else if err := json.Unmarshal(row.ExecutionConfiguration, &projection); err != nil { + return v1.SessionExecutionConfiguration{}, errors.New("invalid stored session execution configuration") + } + sessions.NormalizeExecutionProjection(&projection, uuid.UUID(row.ID.Bytes).String()) + return projection, nil +} + +func (s *Store) MeasuredSessionUsage(ctx context.Context, tenantID, sessionID string) (json.RawMessage, error) { + tenant, err := parseID(tenantID) + if err != nil { + return nil, err + } + id, err := parseID(sessionID) + if err != nil { + return nil, err + } + usage, err := s.units.Queries().SessionMeasuredTokenUsage(ctx, sqlc.SessionMeasuredTokenUsageParams{TenantID: tenant, ID: id}) + if err != nil { + return nil, fmt.Errorf("read measured session usage: %w", err) + } + return usage, nil +} + +func (s *Store) GetManagedSessionArchive(ctx context.Context, tenantID, sessionID string) (sessions.ManagedArchive, error) { + tenant, err := parseID(tenantID) + if err != nil { + return sessions.ManagedArchive{}, err + } + return loadManagedArchive(ctx, s.units.Queries(), tenant, pgunit.PathID(sessionID)) +} + +// loadManagedArchive reads the resource disposal of the tenant's visible +// hosted Session. +func loadManagedArchive(ctx context.Context, q *sqlc.Queries, tenant, session pgtype.UUID) (sessions.ManagedArchive, error) { + row, err := q.GetManagedSessionArchive(ctx, sqlc.GetManagedSessionArchiveParams{TenantID: tenant, ID: session}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.ManagedArchive{}, sessions.ErrNotFound + } + if err != nil { + return sessions.ManagedArchive{}, err + } + if row.EnvironmentType != "openai_hosted" { + return sessions.ManagedArchive{}, sessions.ErrInvalidInput + } + return sessions.ManagedArchive{SessionID: uuid.UUID(row.SessionID.Bytes).String(), EnvironmentID: uuid.UUID(row.EnvironmentID.Bytes).String(), State: row.State}, nil +} + +// LoadSessionActivity adds the Environment, input activity and latest Turn +// projection of GetSession to session within the caller's snapshot. +func LoadSessionActivity(ctx context.Context, q *sqlc.Queries, session sessions.Session) (sessions.Session, error) { + id, _ := parseID(session.ID) + tenant, _ := parseID(session.TenantID) + environment, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: tenant, ID: id}) + if err == nil { + value, err := EnvironmentFromRow(environment.Environment, environment.TenantID, environment.Configuration, nil) + if err != nil { + return session, err + } + session.Environment = &value + session.EnvironmentFailure = environmentFailure(environment.Environment) + state, err := LoadEnvironmentInput(ctx, q, id) + if err != nil { + return session, err + } + session.EnvironmentInputActivity, session.PendingInput = sessions.InputActivity(state) + } else if !errors.Is(err, pgx.ErrNoRows) { + return session, err + } + row, err := q.GetLatestSessionTurn(ctx, id) + if errors.Is(err, pgx.ErrNoRows) { + return session, nil + } + if err != nil { + return session, err + } + turn := TurnFromRow(row) + session.LastTurn = &turn + session.RequiredActions, err = LoadRequiredActions(ctx, q, id, turn) + if err != nil { + return session, err + } + session.Usage, err = q.SessionTokenUsage(ctx, id) + return session, err +} + +// environmentFailure returns the recorded provisioning failure of a failed +// Environment, with its detail sanitized. +func environmentFailure(row sqlc.Environment) *sessions.EnvironmentFailure { + if row.Status != "failed" || !row.FailureReason.Valid || !row.FailedAt.Valid { + return nil + } + failure := &sessions.EnvironmentFailure{Reason: row.FailureReason.String, FailedAt: row.FailedAt.Time} + var detail sessions.ProvisioningFailureDetail + if json.Unmarshal(row.FailureDetail, &detail) == nil { + failure.Detail = sessions.SanitizedProvisioningDetail(detail) + } + return failure +} diff --git a/services/core/internal/persistence/postgres/sessionpg/session_writes.go b/services/core/internal/persistence/postgres/sessionpg/session_writes.go new file mode 100644 index 000000000..d18311a07 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/session_writes.go @@ -0,0 +1,86 @@ +package sessionpg + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var _ sessions.SessionStorage = (*Store)(nil) + +func (s *Store) WithSessionDeletion(ctx context.Context, tenantID, sessionID string, apply func(context.Context, sessions.LockedSession, sessions.SessionDeletionTx) error) error { + tenant, err := parseID(tenantID) + if err != nil { + return err + } + session := pgunit.PathID(sessionID) + return WithSession(ctx, s.units, tenant, session, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + return apply(ctx, locked, &deletionTx{SessionTx: BindSession(q, tenant, session), tenantID: tenantID}) + }) +} + +// deletionTx is one Session deletion. tenantID is the caller's tenant as the +// write audit records it. +type deletionTx struct { + *SessionTx + tenantID string +} + +func (t *deletionTx) ApplyDeletion(ctx context.Context) error { + if err := t.q.DeleteSessionArtifacts(ctx, t.session); err != nil { + return err + } + if err := t.q.ReleaseUnallocatedRuntimePlacement(ctx, t.session); err != nil { + return err + } + return t.q.MarkSessionDeleted(ctx, t.session) +} + +func (t *deletionTx) RecordDeletionAudit(ctx context.Context) error { + return auditpg.RecordWriteAudit(ctx, t.q, t.tenantID, "delete", "session", uuid.UUID(t.session.Bytes).String(), "") +} + +func (s *Store) UpdateSessionMetadata(ctx context.Context, tenantID, sessionID string, encoded []byte) (sessions.Session, error) { + tenant, err := parseID(tenantID) + if err != nil { + return sessions.Session{}, err + } + var session sessions.Session + err = s.units.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + row, err := q.UpdateSessionMetadata(ctx, sqlc.UpdateSessionMetadataParams{TenantID: tenant, ID: pgunit.PathID(sessionID), Metadata: encoded}) + if err != nil { + return err + } + if err := auditpg.RecordWriteAudit(ctx, q, tenantID, "update", "session", uuid.UUID(row.ID.Bytes).String(), ""); err != nil { + return err + } + if session, err = SessionFromRow(row); err != nil { + return err + } + session, err = LoadSessionActivity(ctx, q, session) + return err + }) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.Session{}, sessions.ErrNotFound + } + if err != nil { + return sessions.Session{}, fmt.Errorf("update session metadata: %w", err) + } + return session, nil +} + +func (s *Store) AuditSessionOperation(ctx context.Context, tenantID, sessionID, action string) error { + return s.withPublicSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { + return auditpg.RecordWriteAudit(ctx, q, tenantID, action, "session", uuid.UUID(session.Bytes).String(), "") + }) +} diff --git a/services/core/internal/persistence/postgres/sessionpg/turn_reads.go b/services/core/internal/persistence/postgres/sessionpg/turn_reads.go index 8704de091..9848a68d9 100644 --- a/services/core/internal/persistence/postgres/sessionpg/turn_reads.go +++ b/services/core/internal/persistence/postgres/sessionpg/turn_reads.go @@ -17,7 +17,7 @@ import ( var _ sessions.TurnReader = (*Store)(nil) func (s *Store) GetTurn(ctx context.Context, tenant, session, turn string) (sessions.Turn, error) { - lookup, err := PublicTurnLookup(tenant, session, turn) + lookup, err := publicTurnLookup(tenant, session, turn) if err != nil { return sessions.Turn{}, err } @@ -110,10 +110,10 @@ func ExecutionWorkCursor(after string, connectedDevices []string) (pgtype.UUID, return id, devices, nil } -// PublicTurnLookup resolves the caller's path identifiers of a Turn or a +// publicTurnLookup resolves the caller's path identifiers of a Turn or a // Turn-scoped resource. A malformed tenant is sessions.ErrInvalidInput; a // malformed Session or Turn ID resolves as a missing one. -func PublicTurnLookup(tenant, session, turn string) (sqlc.GetTurnParams, error) { +func publicTurnLookup(tenant, session, turn string) (sqlc.GetTurnParams, error) { id, err := parseID(tenant) return sqlc.GetTurnParams{TenantID: id, SessionID: pgunit.PathID(session), ID: pgunit.PathID(turn)}, err } diff --git a/services/core/internal/runtimehistory/service.go b/services/core/internal/runtimehistory/service.go index 59043201c..2a8be7b60 100644 --- a/services/core/internal/runtimehistory/service.go +++ b/services/core/internal/runtimehistory/service.go @@ -2,8 +2,12 @@ package runtimehistory import ( "context" + "encoding/json" "errors" + "fmt" "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) var ( @@ -13,31 +17,27 @@ var ( ErrUnsupported = errors.New("Runtime history is unsupported for this Session") ) -type ScopeResolver interface { - ResolveRuntimeHistoryScope(context.Context, string, string) (Scope, error) -} - type Reader interface { Capabilities() Capabilities Query(context.Context, Query) (Result, error) } type Service struct { - resolver ScopeResolver + environments sessions.EnvironmentReader reader Reader capabilities Capabilities now func() time.Time } -func NewService(resolver ScopeResolver, reader Reader) (*Service, error) { - if resolver == nil || reader == nil { - return nil, errors.New("Runtime history resolver and reader are required") +func NewService(environments sessions.EnvironmentReader, reader Reader) (*Service, error) { + if environments == nil || reader == nil { + return nil, errors.New("Runtime history Environments and reader are required") } capabilities := reader.Capabilities() if err := capabilities.Validate(); err != nil { return nil, err } - return &Service{resolver: resolver, reader: reader, capabilities: cloneCapabilities(capabilities), now: time.Now}, nil + return &Service{environments: environments, reader: reader, capabilities: cloneCapabilities(capabilities), now: time.Now}, nil } func (s *Service) Capabilities() Capabilities { @@ -52,13 +52,10 @@ func (s *Service) QuerySession(ctx context.Context, tenantID, sessionID string, if !validPublicBoundary(requested.Start) || !validPublicBoundary(requested.End) || !requested.End.After(requested.Start) || requested.End.After(requestNow.Add(time.Second)) || requested.End.Sub(requested.Start) > s.capabilities.MaximumRange || requested.MaxPoints < 2 || requested.MaxPoints > s.capabilities.MaximumPoints { return Response{}, ErrInvalidRange } - scope, err := s.resolver.ResolveRuntimeHistoryScope(ctx, tenantID, sessionID) + scope, err := s.resolveScope(ctx, tenantID, sessionID) if err != nil { return Response{}, err } - if scope.TenantID != tenantID || scope.SessionID != sessionID { - return Response{}, ErrInvalidResult - } if err := scope.validate(); err != nil { return Response{}, ErrInvalidResult } @@ -87,6 +84,30 @@ func (s *Service) QuerySession(ctx context.Context, tenantID, sessionID string, }, nil } +// resolveScope authorizes the Session and returns only durable Core identity. +// It does not resolve a current allocation: retained history may contain +// earlier allocations, and the Reader keeps each durable allocation as one +// continuous series. +func (s *Service) resolveScope(ctx context.Context, tenantID, sessionID string) (Scope, error) { + environment, err := s.environments.GetSessionEnvironment(ctx, tenantID, sessionID) + if err != nil { + return Scope{}, fmt.Errorf("resolve Runtime history Environment: %w", err) + } + if environment.TenantID != tenantID || environment.SessionID != sessionID { + return Scope{}, errors.New("Runtime history Environment does not match resolved ownership") + } + var configuration struct { + Type string `json:"type"` + } + if json.Unmarshal(environment.Configuration, &configuration) != nil || configuration.Type == "" { + return Scope{}, errors.New("invalid stored Runtime history environment configuration") + } + if configuration.Type != "openai_hosted" { + return Scope{}, ErrUnsupported + } + return Scope{TenantID: tenantID, SessionID: sessionID, EnvironmentID: environment.ID}, nil +} + func resolution(duration time.Duration, maxPoints int, minimum time.Duration) time.Duration { divisor := time.Duration(maxPoints) step := duration / divisor diff --git a/services/core/internal/runtimehistory/service_test.go b/services/core/internal/runtimehistory/service_test.go index 63ba36b8e..d56955aa7 100644 --- a/services/core/internal/runtimehistory/service_test.go +++ b/services/core/internal/runtimehistory/service_test.go @@ -5,6 +5,9 @@ import ( "errors" "testing" "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) const ( @@ -14,13 +17,48 @@ const ( allocationID = "44444444-4444-4444-8444-444444444444" ) -type fixedScopeResolver struct { - scope Scope - err error +// fakeEnvironmentReader serves the Session Environment read that scopes a +// query; every other read fails the test. +type fakeEnvironmentReader struct { + t testing.TB + getSessionEnvironment func(context.Context, string, string) (sessions.Environment, error) +} + +func (f *fakeEnvironmentReader) GetSessionEnvironment(ctx context.Context, tenant, session string) (sessions.Environment, error) { + if f.getSessionEnvironment == nil { + f.t.Fatalf("unexpected call to GetSessionEnvironment") + } + return f.getSessionEnvironment(ctx, tenant, session) +} + +func (f *fakeEnvironmentReader) GetEnvironment(context.Context, string, string) (sessions.Environment, error) { + f.t.Fatalf("unexpected call to GetEnvironment") + return sessions.Environment{}, nil +} + +func (f *fakeEnvironmentReader) ListEnvironmentInitializations(context.Context, string) ([]sessions.EnvironmentInitialization, error) { + f.t.Fatalf("unexpected call to ListEnvironmentInitializations") + return nil, nil +} + +func (f *fakeEnvironmentReader) ReadEnvironmentSetup(context.Context, string, string) (environmentconfig.Setup, error) { + f.t.Fatalf("unexpected call to ReadEnvironmentSetup") + return environmentconfig.Setup{}, nil } -func (r fixedScopeResolver) ResolveRuntimeHistoryScope(context.Context, string, string) (Scope, error) { - return r.scope, r.err +func (f *fakeEnvironmentReader) ReadInitialEnvironmentFile(context.Context, string, string, int) (environmentconfig.InitialFileMetadata, []byte, error) { + f.t.Fatalf("unexpected call to ReadInitialEnvironmentFile") + return environmentconfig.InitialFileMetadata{}, nil, nil +} + +// sessionEnvironment reads environment and err as the Session's Environment. +func sessionEnvironment(t testing.TB, environment sessions.Environment, err error) *fakeEnvironmentReader { + return &fakeEnvironmentReader{t: t, getSessionEnvironment: func(context.Context, string, string) (sessions.Environment, error) { return environment, err }} +} + +// hosted reads the hosted Environment of the Session the tests query. +func hosted(t testing.TB) *fakeEnvironmentReader { + return sessionEnvironment(t, sessions.Environment{ID: environmentID, TenantID: tenantID, SessionID: sessionID, Configuration: []byte(`{"type":"openai_hosted"}`)}, nil) } type fakeReader struct { @@ -69,7 +107,7 @@ func TestServiceAuthorizesAndBoundsBackendQuery(t *testing.T) { CPUUtilizationRatio: &ratio, CPUCapacityCores: &capacity, MemoryUsageBytes: &memory, MemoryLimitBytes: &limit, }}, }}} - service, err := NewService(fixedScopeResolver{scope: scope}, reader) + service, err := NewService(hosted(t), reader) if err != nil { t.Fatal(err) } @@ -94,7 +132,7 @@ func TestServiceAuthorizesAndBoundsBackendQuery(t *testing.T) { func TestServiceNeverQueriesBeforeOwnershipResolution(t *testing.T) { reader := &fakeReader{capabilities: capabilities()} denied := errors.New("not found") - service, err := NewService(fixedScopeResolver{err: denied}, reader) + service, err := NewService(sessionEnvironment(t, sessions.Environment{}, denied), reader) if err != nil { t.Fatal(err) } @@ -108,9 +146,8 @@ func TestServiceNeverQueriesBeforeOwnershipResolution(t *testing.T) { func TestServiceValidatesResultAgainstTimeAfterReaderReturns(t *testing.T) { requestNow := time.Date(2026, 9, 23, 8, 0, 0, 0, time.UTC) resultNow := requestNow.Add(2 * time.Second) - scope := Scope{TenantID: tenantID, SessionID: sessionID, EnvironmentID: environmentID} reader := &fakeReader{capabilities: capabilities(), result: Result{GeneratedAt: resultNow}} - service, err := NewService(fixedScopeResolver{scope: scope}, reader) + service, err := NewService(hosted(t), reader) if err != nil { t.Fatal(err) } @@ -136,8 +173,7 @@ func TestServiceValidatesResultAgainstTimeAfterReaderReturns(t *testing.T) { func TestServiceRejectsInvalidRangeAndResolverIdentity(t *testing.T) { reader := &fakeReader{capabilities: capabilities()} - scope := Scope{TenantID: tenantID, SessionID: sessionID, EnvironmentID: environmentID} - service, err := NewService(fixedScopeResolver{scope: scope}, reader) + service, err := NewService(hosted(t), reader) if err != nil { t.Fatal(err) } @@ -155,9 +191,21 @@ func TestServiceRejectsInvalidRangeAndResolverIdentity(t *testing.T) { if len(reader.queries) != 0 { t.Fatalf("invalid range reached reader: %+v", reader.queries) } - service.resolver = fixedScopeResolver{scope: Scope{TenantID: "55555555-5555-4555-8555-555555555555", SessionID: sessionID, EnvironmentID: environmentID}} - if _, err := service.QuerySession(t.Context(), tenantID, sessionID, Range{Start: now.Add(-time.Hour), End: now, MaxPoints: 60}); err == nil || len(reader.queries) != 0 { - t.Fatal("mismatched resolver identity reached reader") + for _, environment := range []sessions.Environment{ + {ID: environmentID, TenantID: tenantID, SessionID: sessionID, Configuration: []byte(`{"type":"self_hosted"}`)}, + {ID: environmentID, TenantID: "55555555-5555-4555-8555-555555555555", SessionID: sessionID, Configuration: []byte(`{"type":"openai_hosted"}`)}, + {ID: environmentID, TenantID: tenantID, SessionID: "66666666-6666-4666-8666-666666666666", Configuration: []byte(`{"type":"openai_hosted"}`)}, + {ID: environmentID, TenantID: tenantID, SessionID: sessionID, Configuration: []byte(`{"type":`)}, + {ID: "environment", TenantID: tenantID, SessionID: sessionID, Configuration: []byte(`{"type":"openai_hosted"}`)}, + } { + service.environments = sessionEnvironment(t, environment, nil) + if _, err := service.QuerySession(t.Context(), tenantID, sessionID, Range{Start: now.Add(-time.Hour), End: now, MaxPoints: 60}); err == nil || len(reader.queries) != 0 { + t.Fatalf("unsafe Runtime history Environment reached reader: %+v", environment) + } + } + service.environments = sessionEnvironment(t, sessions.Environment{ID: environmentID, TenantID: tenantID, SessionID: sessionID, Configuration: []byte(`{"type":"self_hosted"}`)}, nil) + if _, err := service.QuerySession(t.Context(), tenantID, sessionID, Range{Start: now.Add(-time.Hour), End: now, MaxPoints: 60}); !errors.Is(err, ErrUnsupported) { + t.Fatalf("self-hosted history was not unsupported: %v", err) } } @@ -266,7 +314,7 @@ func TestServiceRejectsMalformedBackendResults(t *testing.T) { result.Series[0].Points = []Point{{Start: start, End: start.Add(time.Minute)}} mutate(&result) reader := &fakeReader{capabilities: capabilities(), result: result} - service, err := NewService(fixedScopeResolver{scope: scope}, reader) + service, err := NewService(hosted(t), reader) if err != nil { t.Fatal(err) } @@ -288,7 +336,7 @@ func TestServiceRejectsUnsafeCapabilities(t *testing.T) { value := base value.Metrics = append([]Metric(nil), base.Metrics...) mutate(&value) - if _, err := NewService(fixedScopeResolver{}, &fakeReader{capabilities: value}); err == nil { + if _, err := NewService(hosted(t), &fakeReader{capabilities: value}); err == nil { t.Fatalf("unsafe capabilities accepted: %+v", value) } } diff --git a/services/core/internal/runtimehistory/storeresolver/resolver.go b/services/core/internal/runtimehistory/storeresolver/resolver.go deleted file mode 100644 index 040de4bda..000000000 --- a/services/core/internal/runtimehistory/storeresolver/resolver.go +++ /dev/null @@ -1,48 +0,0 @@ -package storeresolver - -import ( - "context" - "encoding/json" - "errors" - "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimehistory" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -type environmentStore interface { - GetSessionEnvironment(context.Context, string, string) (sessions.Environment, error) -} - -type Resolver struct{ store environmentStore } - -func NewResolver(value environmentStore) (*Resolver, error) { - if value == nil { - return nil, errors.New("Runtime history store is required") - } - return &Resolver{store: value}, nil -} - -// ResolveRuntimeHistoryScope authorizes the Session through the Core store and -// returns only durable Core identity. It intentionally does not resolve a -// current allocation: retained history may contain earlier allocations. The -// Reader keeps each durable allocation as one continuous series. -func (r *Resolver) ResolveRuntimeHistoryScope(ctx context.Context, tenantID, sessionID string) (runtimehistory.Scope, error) { - environment, err := r.store.GetSessionEnvironment(ctx, tenantID, sessionID) - if err != nil { - return runtimehistory.Scope{}, fmt.Errorf("resolve Runtime history Environment: %w", err) - } - if environment.TenantID != tenantID || environment.SessionID != sessionID { - return runtimehistory.Scope{}, errors.New("Runtime history Environment does not match resolved ownership") - } - var configuration struct { - Type string `json:"type"` - } - if json.Unmarshal(environment.Configuration, &configuration) != nil || configuration.Type == "" { - return runtimehistory.Scope{}, errors.New("invalid stored Runtime history environment configuration") - } - if configuration.Type != "openai_hosted" { - return runtimehistory.Scope{}, runtimehistory.ErrUnsupported - } - return runtimehistory.Scope{TenantID: tenantID, SessionID: sessionID, EnvironmentID: environment.ID}, nil -} diff --git a/services/core/internal/runtimehistory/storeresolver/resolver_test.go b/services/core/internal/runtimehistory/storeresolver/resolver_test.go deleted file mode 100644 index 83cf2939f..000000000 --- a/services/core/internal/runtimehistory/storeresolver/resolver_test.go +++ /dev/null @@ -1,74 +0,0 @@ -package storeresolver - -import ( - "context" - "errors" - "testing" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimehistory" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -type resolverStore struct { - environment sessions.Environment - err error - calls int -} - -func (s *resolverStore) GetSessionEnvironment(context.Context, string, string) (sessions.Environment, error) { - s.calls++ - return s.environment, s.err -} - -func TestResolverAuthorizesManagedSessionWithoutSelectingCurrentAllocation(t *testing.T) { - backend := &resolverStore{environment: sessions.Environment{ - ID: environmentID, TenantID: tenantID, SessionID: sessionID, - Configuration: []byte(`{"type":"openai_hosted"}`), - }} - resolver, err := NewResolver(backend) - if err != nil { - t.Fatal(err) - } - scope, err := resolver.ResolveRuntimeHistoryScope(t.Context(), tenantID, sessionID) - if err != nil { - t.Fatal(err) - } - if scope != (runtimehistory.Scope{TenantID: tenantID, SessionID: sessionID, EnvironmentID: environmentID}) || backend.calls != 1 { - t.Fatalf("unexpected Runtime history scope: %+v calls=%d", scope, backend.calls) - } -} - -func TestResolverPreservesTenantScopedNotFound(t *testing.T) { - backend := &resolverStore{err: sessions.ErrNotFound} - resolver, err := NewResolver(backend) - if err != nil { - t.Fatal(err) - } - _, err = resolver.ResolveRuntimeHistoryScope(t.Context(), tenantID, sessionID) - if !errors.Is(err, sessions.ErrNotFound) { - t.Fatalf("tenant-scoped not found was not preserved: %v", err) - } -} - -func TestResolverRejectsUnsupportedOrMismatchedEnvironment(t *testing.T) { - for _, environment := range []sessions.Environment{ - {ID: environmentID, TenantID: tenantID, SessionID: sessionID, Configuration: []byte(`{"type":"self_hosted"}`)}, - {ID: environmentID, TenantID: "55555555-5555-4555-8555-555555555555", SessionID: sessionID, Configuration: []byte(`{"type":"openai_hosted"}`)}, - {ID: environmentID, TenantID: tenantID, SessionID: "66666666-6666-4666-8666-666666666666", Configuration: []byte(`{"type":"openai_hosted"}`)}, - {ID: environmentID, TenantID: tenantID, SessionID: sessionID, Configuration: []byte(`{"type":`)}, - } { - resolver, err := NewResolver(&resolverStore{environment: environment}) - if err != nil { - t.Fatal(err) - } - if _, err := resolver.ResolveRuntimeHistoryScope(t.Context(), tenantID, sessionID); err == nil { - t.Fatalf("unsafe Runtime history Environment accepted: %+v", environment) - } - } -} - -const ( - tenantID = "11111111-1111-4111-8111-111111111111" - sessionID = "22222222-2222-4222-8222-222222222222" - environmentID = "33333333-3333-4333-8333-333333333333" -) diff --git a/services/core/internal/sandbox/providers/configuration_flow_test.go b/services/core/internal/sandbox/providers/configuration_flow_test.go index b82a51051..2fd9a314c 100644 --- a/services/core/internal/sandbox/providers/configuration_flow_test.go +++ b/services/core/internal/sandbox/providers/configuration_flow_test.go @@ -133,15 +133,16 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { EnvironmentTemplates: struct{ api.EnvironmentTemplates }{}, EnvironmentTemplatesReader: struct{ api.EnvironmentTemplatesReader }{}, Skills: struct{ api.Skills }{}, SkillsReader: struct{ api.SkillsReader }{}, Agents: struct{ api.Agents }{}, AgentsReader: struct{ api.AgentsReader }{}, - Sessions: s, + Sessions: struct{ api.Sessions }{}, + SessionsReader: struct{ api.SessionsReader }{}, SessionCreation: s, - SessionEvents: s, + SessionEvents: struct{ api.SessionEvents }{}, Turns: struct{ api.Turns }{}, Items: struct{ api.Items }{}, Subagents: struct{ api.Subagents }{}, Artifacts: struct{ api.Artifacts }{}, ArtifactsReader: struct{ api.ArtifactsReader }{}, - SessionAdmin: s, + SessionAdmin: struct{ api.SessionAdmin }{}, Environments: struct{ api.Environments }{}, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, Admin: s, AdminRuntimeTargets: struct{ api.AdminRuntimeTargets }{}, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, ExecutorConnections: struct{ api.ExecutorConnections }{}, Metrics: struct{ api.Metrics }{}, RuntimeObservations: struct{ api.RuntimeObservations }{}, RuntimeHistory: struct{ api.RuntimeHistory }{}, @@ -149,7 +150,7 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { ExecutorURL: "wss://core.example/api/v1/agent-daemon/ws", SessionAdmission: s, InputAdmission: s, - SessionArchive: s, + SessionArchive: struct{ api.SessionArchive }{}, Workspaces: struct{ api.EnvironmentWorkspaces }{}, }, Sandboxes: &api.Sandboxes{Deployment: service, NodeAllocations: deploymentpg.New(pgunit.NewPool(pool), nil), DeploymentChanges: leaseSetup{t: t, changes: changes, installation: installation}, diff --git a/services/core/internal/sessions/artifacts_test.go b/services/core/internal/sessions/artifacts_test.go index 7132974d7..a31e5b326 100644 --- a/services/core/internal/sessions/artifacts_test.go +++ b/services/core/internal/sessions/artifacts_test.go @@ -257,6 +257,21 @@ func (s *fakeArtifactStorage) WithEnvironmentExecutorCredentials(context.Context return nil } +func (s *fakeArtifactStorage) WithSessionDeletion(context.Context, string, string, func(context.Context, LockedSession, SessionDeletionTx) error) error { + s.t.Fatal("unexpected call to WithSessionDeletion") + return nil +} + +func (s *fakeArtifactStorage) UpdateSessionMetadata(context.Context, string, string, []byte) (Session, error) { + s.t.Fatal("unexpected call to UpdateSessionMetadata") + return Session{}, nil +} + +func (s *fakeArtifactStorage) AuditSessionOperation(context.Context, string, string, string) error { + s.t.Fatal("unexpected call to AuditSessionOperation") + 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 20f8676e5..b8297e341 100644 --- a/services/core/internal/sessions/devices_test.go +++ b/services/core/internal/sessions/devices_test.go @@ -40,6 +40,11 @@ type fakeStorage struct { // credentials is the transaction the executor credential With methods // apply in, with locked. credentials *fakeTx + // deletion is the transaction WithSessionDeletion applies in, with + // locked. + deletion *fakeTx + updateMetadata func(encoded string) (Session, error) + auditOperation func() error } func (s *fakeStorage) record(name string, set bool, detail ...string) { @@ -75,6 +80,21 @@ func (s *fakeStorage) WithEnrollment(ctx context.Context, environment, credentia return apply(ctx, s.enrollment, s.environment, s.locked) } +func (s *fakeStorage) WithSessionDeletion(ctx context.Context, tenant, session string, apply func(context.Context, LockedSession, SessionDeletionTx) error) error { + s.record("WithSessionDeletion", s.deletion != nil, tenant, session) + return apply(ctx, s.locked, s.deletion) +} + +func (s *fakeStorage) UpdateSessionMetadata(_ context.Context, tenant, session string, encoded []byte) (Session, error) { + s.record("UpdateSessionMetadata", s.updateMetadata != nil, tenant, session, string(encoded)) + return s.updateMetadata(string(encoded)) +} + +func (s *fakeStorage) AuditSessionOperation(_ context.Context, tenant, session, action string) error { + s.record("AuditSessionOperation", s.auditOperation != nil, tenant, session, action) + return s.auditOperation() +} + func (s *fakeStorage) WithArtifactStaging(context.Context, ArtifactStagingKey, func(context.Context, ArtifactStagingTx) error) error { s.t.Fatal("unexpected call to WithArtifactStaging") return nil diff --git a/services/core/internal/sessions/execution_configuration.go b/services/core/internal/sessions/execution_configuration.go new file mode 100644 index 000000000..6065dbe97 --- /dev/null +++ b/services/core/internal/sessions/execution_configuration.go @@ -0,0 +1,44 @@ +package sessions + +import ( + "encoding/json" + "errors" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" +) + +// NormalizeExecutionProjection completes a Session's frozen execution +// configuration for sessionID as it is stored and read: the object and schema +// version, an unknown harness configuration source, the resolved harness +// configuration, and the provider view only while it is available. A +// deployment provider frozen without its safe view stays redacted. +func NormalizeExecutionProjection(projection *v1.SessionExecutionConfiguration, sessionID string) { + projection.Object = "agent.session.execution_configuration" + projection.SchemaVersion = 1 + if projection.HarnessConfig.Source == "" { + projection.HarnessConfig.Source = "unknown" + } + projection.HarnessConfig.Value = v1.ResolvedHarnessConfig(projection.HarnessConfig.Value) + projection.SessionID = sessionID + if projection.ModelProvider.Source == "deployment" && (projection.ModelProvider.Status != "available" || projection.ModelProvider.Configuration == nil) { + // Sessions created before deployment defaults moved into Core stay redacted. + projection.ModelProvider.Status = "redacted" + projection.ModelProvider.Configuration = nil + } else if projection.ModelProvider.Status != "available" { + projection.ModelProvider.Configuration = nil + } +} + +// ExecutionModel returns the model a stored Session configuration selects, +// nil when it selects none. +func ExecutionModel(configuration []byte) (*string, error) { + var config struct { + Agent struct { + Model *string `json:"model"` + } `json:"agent"` + } + if err := json.Unmarshal(configuration, &config); err != nil { + return nil, errors.New("invalid stored Session model configuration") + } + return config.Agent.Model, nil +} diff --git a/services/core/internal/sessions/reader.go b/services/core/internal/sessions/reader.go index 32dd253bd..6da1497a6 100644 --- a/services/core/internal/sessions/reader.go +++ b/services/core/internal/sessions/reader.go @@ -1,6 +1,6 @@ package sessions -// Reader reads a Session's resources, Environment and devices, and the +// Reader reads Sessions and their resources, Environment and devices, and the // administrator's cross-Project Session views, one family per line. Callers // use it directly; no use case forwards a read. type Reader interface { @@ -11,6 +11,7 @@ type Reader interface { ExecutorCredentialReader ItemReader ModelExecutionReader + SessionReader SubagentReader TurnReader } diff --git a/services/core/internal/sessions/reads.go b/services/core/internal/sessions/reads.go index f8fb70a74..883dd24b8 100644 --- a/services/core/internal/sessions/reads.go +++ b/services/core/internal/sessions/reads.go @@ -1,6 +1,53 @@ package sessions -import "time" +import ( + "context" + "encoding/json" + "time" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" +) + +// SessionReader reads Sessions, their public change journal and their +// diagnostics. +type SessionReader interface { + // GetSession returns the tenant's visible Session with its Environment, + // input activity and latest Turn, read from one snapshot, or ErrNotFound. + GetSession(ctx context.Context, tenantID, sessionID string) (Session, error) + // ListSessions returns one page of the tenant's visible Sessions in + // creation-time and ID order; a non-nil agentID keeps that root Agent's + // Sessions. The cursor is the last returned Session ID and must belong to + // the tenant; it grants no additional access. A limit outside 1..100 is + // ErrInvalidInput. + ListSessions(ctx context.Context, tenantID, cursor string, limit int, ascending bool, agentID *string) (Page, error) + // SessionStreamSnapshot reads the projection GetSession returns and the + // committed Session event cursor from the same snapshot. + SessionStreamSnapshot(ctx context.Context, tenantID, sessionID string) (Session, int64, error) + // SessionEventCursor returns the sequence of the visible Session's latest + // public change, or ErrNotFound. + SessionEventCursor(ctx context.Context, tenantID, sessionID string) (int64, error) + // ListSessionEvents returns the retained public changes after the + // sequence, in order. A negative after is ErrInvalidInput; a pruned change + // in between is ErrStreamGap. + ListSessionEvents(ctx context.Context, tenantID, sessionID string, after int64) ([]SessionChange, error) + // GetTurnDiagnosticsSnapshot reads a root Turn and the timing of up to + // 1000 of its root Items with the Session in one read-only snapshot. + GetTurnDiagnosticsSnapshot(ctx context.Context, tenantID, sessionID, turnID string) (TurnDiagnosticsSnapshot, error) + // GetSessionExecutionConfiguration reads only the safe committed + // configuration. It never loads provider ciphertext, current defaults or + // runtime health. + GetSessionExecutionConfiguration(ctx context.Context, tenantID, sessionID string) (v1.SessionExecutionConfiguration, error) + // MeasuredSessionUsage returns Core-internal measured usage for Runtime + // telemetry: the sum of every recorded root Turn snapshot, active Turns + // included, and null only when nothing is recorded. Public Session usage + // keeps the official rule. A missing Session reads as null, so callers + // resolve the Session first. + MeasuredSessionUsage(ctx context.Context, tenantID, sessionID string) (json.RawMessage, error) + // GetManagedSessionArchive reads the resource disposal of the tenant's + // hosted Session and never contacts compute. A missing Session is + // ErrNotFound; one that is not hosted is ErrInvalidInput. + GetManagedSessionArchive(ctx context.Context, tenantID, sessionID string) (ManagedArchive, error) +} // ItemDiagnosticTiming records Core database receipt and settlement, never native // execution duration. Historical terminal Items can have unknown settlement. diff --git a/services/core/internal/sessions/service.go b/services/core/internal/sessions/service.go index 4ac82fdc9..8612a15e7 100644 --- a/services/core/internal/sessions/service.go +++ b/services/core/internal/sessions/service.go @@ -19,4 +19,5 @@ type Storage interface { ArtifactStorage DeviceStorage ExecutorCredentialStorage + SessionStorage } diff --git a/services/core/internal/sessions/session.go b/services/core/internal/sessions/session.go index 61ded8a8d..8e922604d 100644 --- a/services/core/internal/sessions/session.go +++ b/services/core/internal/sessions/session.go @@ -1,7 +1,9 @@ package sessions import ( + "context" "encoding/json" + "fmt" "regexp" "time" @@ -10,6 +12,7 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/metadata" ) var enginePattern = regexp.MustCompile(`^[a-z][a-z0-9_-]{0,63}$`) @@ -73,3 +76,117 @@ type Creation struct { } func ValidEngine(engine string) bool { return enginePattern.MatchString(engine) } + +// SessionStorage persists the public Session writes. +type SessionStorage interface { + // WithSessionDeletion runs apply in one Session transaction of the + // tenant's Session, publicly deleted or not, with what the Session lock + // shows. It prunes the Session's journal after apply and commits only + // when both succeed. A malformed tenant is ErrInvalidInput; a malformed or + // missing Session is ErrNotFound. + WithSessionDeletion(ctx context.Context, tenantID, sessionID string, apply func(context.Context, LockedSession, SessionDeletionTx) error) error + // UpdateSessionMetadata replaces the visible Session's metadata with the + // encoded object, records the write audit and reads the Session as + // GetSession does, in one transaction, or returns ErrNotFound. + UpdateSessionMetadata(ctx context.Context, tenantID, sessionID string, encoded []byte) (Session, error) + // AuditSessionOperation records the write audit of action on the visible + // Session under its lock, or returns ErrNotFound. + AuditSessionOperation(ctx context.Context, tenantID, sessionID, action string) error +} + +// SessionDeletionTx is the Session transaction DeleteSession runs in. +type SessionDeletionTx interface { + ActiveTurnTx + // LoadEnvironmentInput reads the Session's latest Environment input + // reservation that no Turn has admitted or superseded, nil when there is + // none. + LoadEnvironmentInput(ctx context.Context) (*EnvironmentInputState, error) + // ApplyDeletion deletes the Session's Artifacts, releases the node + // placement of its Environments that have no allocation and removes the + // Session from public access. Its row stays so that execution can settle. + ApplyDeletion(ctx context.Context) error + // RecordDeletionAudit records the write audit of the deletion. + RecordDeletionAudit(ctx context.Context) error +} + +// settled reports whether a Session may be deleted: it has no queued, +// in-progress or waiting Turn, which includes pending required actions and +// function results, and no pending input reservation, such as queued later +// input, self-hosted input awaiting a connection, or hosted initial input while +// provisioning. Terminal idle and failed Sessions are settled. +func settled(activeTurn bool, input *EnvironmentInputState) bool { + _, pending := InputActivity(input) + return !activeTurn && !pending +} + +// DeleteSessionCommand deletes one Session. +type DeleteSessionCommand struct { + TenantID string + SessionID string +} + +// DeleteSession removes public access to a settled Session while keeping the +// state execution needs to settle. It decides under the Session lock that also +// orders Turn and input admission, so a concurrent admission either commits +// first and is rejected here with ErrNotIdle, or observes the deletion. A +// repeated deletion only records the write audit; foreign and missing Sessions +// are ErrNotFound. +func (s *Service) DeleteSession(ctx context.Context, command DeleteSessionCommand) error { + return s.storage.WithSessionDeletion(ctx, command.TenantID, command.SessionID, func(ctx context.Context, locked LockedSession, tx SessionDeletionTx) error { + if !locked.Deleted { + _, active, err := tx.LoadActiveTurn(ctx) + if err != nil { + return err + } + var input *EnvironmentInputState + if !active { + if input, err = tx.LoadEnvironmentInput(ctx); err != nil { + return err + } + } + if !settled(active, input) { + return ErrNotIdle + } + if err := tx.ApplyDeletion(ctx); err != nil { + return err + } + } + return tx.RecordDeletionAudit(ctx) + }) +} + +// UpdateSessionMetadataCommand replaces a Session's metadata. +type UpdateSessionMetadataCommand struct { + TenantID string + SessionID string + Metadata map[string]string +} + +// UpdateSessionMetadata replaces the visible Session's metadata and returns +// the Session. Invalid metadata is ErrInvalidInput. +func (s *Service) UpdateSessionMetadata(ctx context.Context, command UpdateSessionMetadataCommand) (Session, error) { + encoded, err := metadata.Encode(command.Metadata) + if err != nil { + return Session{}, fmt.Errorf("%w: %w", ErrInvalidInput, err) + } + return s.storage.UpdateSessionMetadata(ctx, command.TenantID, command.SessionID, encoded) +} + +// AuditSessionOperationCommand records an authorized public operation that +// changes no Session: a no-op or a creation replay. +type AuditSessionOperationCommand struct { + TenantID string + SessionID string + // Action is create or send_events. + Action string +} + +// AuditSessionOperation records the write audit of a public no-op or creation +// replay on the visible Session. It cannot create ownership or admit execution +// work. Another action is ErrInvalidInput. +func (s *Service) AuditSessionOperation(ctx context.Context, command AuditSessionOperationCommand) error { + if command.Action != "create" && command.Action != "send_events" { + return ErrInvalidInput + } + return s.storage.AuditSessionOperation(ctx, command.TenantID, command.SessionID, command.Action) +} diff --git a/services/core/internal/sessions/session_test.go b/services/core/internal/sessions/session_test.go new file mode 100644 index 000000000..120ae30cd --- /dev/null +++ b/services/core/internal/sessions/session_test.go @@ -0,0 +1,102 @@ +package sessions + +import ( + "errors" + "strings" + "testing" +) + +func TestSettled(t *testing.T) { + for name, test := range map[string]struct { + active bool + input *EnvironmentInputState + want bool + }{ + "idle": {false, nil, true}, + "active Turn": {true, nil, false}, + "pending input": {false, &EnvironmentInputState{State: EnvironmentInputPending}, false}, + "pending hosted initial": {false, &EnvironmentInputState{State: EnvironmentInputPending, Initial: true, EnvironmentType: "openai_hosted"}, false}, + "failed input": {false, &EnvironmentInputState{State: EnvironmentInputFailed}, true}, + "expired input": {false, &EnvironmentInputState{State: EnvironmentInputExpired}, true}, + "active Turn without input": {true, &EnvironmentInputState{State: EnvironmentInputFailed}, false}, + } { + if got := settled(test.active, test.input); got != test.want { + t.Errorf("%s: settled = %v, want %v", name, got, test.want) + } + } +} + +func TestDeleteSession(t *testing.T) { + errStorage := errors.New("storage") + pending := &EnvironmentInputState{State: EnvironmentInputPending} + for _, test := range []struct { + name string + locked LockedSession + tx fakeTx + want error + calls []string + }{ + {"settled", LockedSession{}, fakeTx{loadActiveTurn: activeTurn(nil), loadEnvironmentInput: returns[*EnvironmentInputState](nil), applyDeletion: done, recordDeletionAudit: done}, nil, + []string{"LoadActiveTurn", "LoadEnvironmentInput", "ApplyDeletion", "RecordDeletionAudit"}}, + {"repeated deletion only audits", LockedSession{Deleted: true}, fakeTx{recordDeletionAudit: done}, nil, []string{"RecordDeletionAudit"}}, + {"active Turn", LockedSession{}, fakeTx{loadActiveTurn: activeTurn(&Turn{ID: "turn"})}, ErrNotIdle, []string{"LoadActiveTurn"}}, + {"pending input", LockedSession{}, fakeTx{loadActiveTurn: activeTurn(nil), loadEnvironmentInput: returns(pending)}, ErrNotIdle, []string{"LoadActiveTurn", "LoadEnvironmentInput"}}, + {"failed deletion", LockedSession{}, fakeTx{loadActiveTurn: activeTurn(nil), loadEnvironmentInput: returns[*EnvironmentInputState](nil), applyDeletion: func() error { return errStorage }}, errStorage, + []string{"LoadActiveTurn", "LoadEnvironmentInput", "ApplyDeletion"}}, + } { + t.Run(test.name, func(t *testing.T) { + tx := test.tx + tx.t = t + storage := &fakeStorage{t: t, deletion: &tx, locked: test.locked} + service, err := NewService(storage) + if err != nil { + t.Fatal(err) + } + err = service.DeleteSession(t.Context(), DeleteSessionCommand{TenantID: "tenant", SessionID: "session"}) + if test.want == nil && err != nil || test.want != nil && !errors.Is(err, test.want) { + t.Fatalf("got %v, want %v", err, test.want) + } + if got := strings.Join(storage.calls, "\n"); got != "WithSessionDeletion tenant session" { + t.Fatalf("storage calls %q", got) + } + if strings.Join(tx.calls, "\n") != strings.Join(test.calls, "\n") { + t.Fatalf("calls %q, want %q", tx.calls, test.calls) + } + }) + } +} + +func TestUpdateSessionMetadata(t *testing.T) { + storage := &fakeStorage{t: t, updateMetadata: func(string) (Session, error) { return Session{ID: "session"}, nil }} + service, err := NewService(storage) + if err != nil { + t.Fatal(err) + } + session, err := service.UpdateSessionMetadata(t.Context(), UpdateSessionMetadataCommand{TenantID: "tenant", SessionID: "session", Metadata: map[string]string{"k": "v"}}) + if err != nil || session.ID != "session" || strings.Join(storage.calls, "\n") != `UpdateSessionMetadata tenant session {"k":"v"}` { + t.Fatal(session, err, storage.calls) + } + invalid := map[string]string{"k": strings.Repeat("v", 64*1024)} + if _, err := service.UpdateSessionMetadata(t.Context(), UpdateSessionMetadataCommand{TenantID: "tenant", SessionID: "session", Metadata: invalid}); !errors.Is(err, ErrInvalidInput) || len(storage.calls) != 1 { + t.Fatal("invalid metadata reached storage", err, storage.calls) + } +} + +func TestAuditSessionOperation(t *testing.T) { + storage := &fakeStorage{t: t, auditOperation: done} + service, err := NewService(storage) + if err != nil { + t.Fatal(err) + } + for _, action := range []string{"create", "send_events"} { + if err := service.AuditSessionOperation(t.Context(), AuditSessionOperationCommand{TenantID: "tenant", SessionID: "session", Action: action}); err != nil { + t.Fatal(action, err) + } + } + if err := service.AuditSessionOperation(t.Context(), AuditSessionOperationCommand{TenantID: "tenant", SessionID: "session", Action: "delete"}); !errors.Is(err, ErrInvalidInput) { + t.Fatal("unsupported action", err) + } + if got := strings.Join(storage.calls, "\n"); got != "AuditSessionOperation tenant session create\nAuditSessionOperation tenant session send_events" { + t.Fatalf("calls %q", got) + } +} diff --git a/services/core/internal/sessions/transaction_test.go b/services/core/internal/sessions/transaction_test.go index 6cc0809bd..318013ab3 100644 --- a/services/core/internal/sessions/transaction_test.go +++ b/services/core/internal/sessions/transaction_test.go @@ -94,6 +94,9 @@ type fakeTx struct { hasUnappliedInputs func() (bool, error) rememberNativeSession func() error beginArtifactCapture func() error + + applyDeletion func() error + recordDeletionAudit func() error } var ( @@ -111,6 +114,7 @@ var ( _ TurnTx = (*fakeTx)(nil) _ EnvironmentExecutorCredentialTx = (*fakeTx)(nil) + _ SessionDeletionTx = (*fakeTx)(nil) ) func (f *fakeTx) record(name string, set bool, detail ...string) { @@ -121,6 +125,16 @@ func (f *fakeTx) record(name string, set bool, detail ...string) { f.calls = append(f.calls, strings.Join(append([]string{name}, detail...), " ")) } +func (f *fakeTx) ApplyDeletion(context.Context) error { + f.record("ApplyDeletion", f.applyDeletion != nil) + return f.applyDeletion() +} + +func (f *fakeTx) RecordDeletionAudit(context.Context) error { + f.record("RecordDeletionAudit", f.recordDeletionAudit != nil) + return f.recordDeletionAudit() +} + func (f *fakeTx) LoadUsage(context.Context) (json.RawMessage, error) { f.record("LoadUsage", f.loadUsage != nil) return f.loadUsage() diff --git a/services/core/internal/store/admin_delete_audit_test.go b/services/core/internal/store/admin_delete_audit_test.go index 8650556f2..4c4acba0f 100644 --- a/services/core/internal/store/admin_delete_audit_test.go +++ b/services/core/internal/store/admin_delete_audit_test.go @@ -161,7 +161,7 @@ func prepareAdminHistoryDelete(t *testing.T, s *Store, name string) (string, res t.Fatal("artifact fixture", err) } verify := func() { - if _, err := s.GetSession(t.Context(), tenant, session); err != nil { + if _, err := sessionAdapter(s).GetSession(t.Context(), tenant, session); err != nil { t.Fatal("Session was not restored", err) } for _, artifact := range page.Artifacts { @@ -181,7 +181,9 @@ func prepareAdminHistoryDelete(t *testing.T, s *Store, name string) (string, res } } if name == "session_delete" { - return tenant, resourceAuditMutation{action: "delete", kind: "session", run: func(ctx context.Context) (string, error) { return session, s.DeleteSession(ctx, tenant, session) }}, verify, 2 + return tenant, resourceAuditMutation{action: "delete", kind: "session", run: func(ctx context.Context) (string, error) { + return session, sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session}) + }}, verify, 2 } artifact := page.Artifacts[0] return tenant, resourceAuditMutation{action: "delete", kind: "artifact", parent: session, run: func(ctx context.Context) (string, error) { @@ -194,7 +196,7 @@ func assertAdminDeletedResource(t *testing.T, s *Store, tenant string, mutation var err error switch mutation.kind { case "session": - _, err = s.GetSession(t.Context(), tenant, id) + _, err = sessionAdapter(s).GetSession(t.Context(), tenant, id) case "artifact": _, err = sessionAdapter(s).GetSessionArtifact(t.Context(), tenant, mutation.parent, id) default: diff --git a/services/core/internal/store/admin_session_archive.go b/services/core/internal/store/admin_session_archive.go deleted file mode 100644 index 2ee7157cb..000000000 --- a/services/core/internal/store/admin_session_archive.go +++ /dev/null @@ -1,161 +0,0 @@ -package store - -import ( - "context" - "errors" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/adminaudit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -// ArchiveManagedSession ends a managed Environment's lifetime while keeping its -// public Session, history and persisted files. The lifecycle owner performs the -// external cleanup; only its existing confirmation can release an allocation. -func (s *Store) ArchiveManagedSession(ctx context.Context, tenantID, sessionID string, expectedGeneration uint64) (sessions.ManagedArchive, error) { - return s.archiveManagedSession(ctx, tenantID, sessionID, expectedGeneration, nil) -} - -// A reset instance is identified by its persisted request time as well as its -// generation, preventing a cancelled clear's candidates from affecting its successor. -func (s *Store) ArchiveSandboxResetSession(ctx context.Context, tenantID, sessionID string, generation uint64, requestedAt time.Time) (sessions.ManagedArchive, error) { - return s.archiveManagedSession(ctx, tenantID, sessionID, generation, &requestedAt) -} - -var ErrSandboxResetSessionBusy = errors.New("the hosted Session is busy") - -func (s *Store) archiveManagedSession(ctx context.Context, tenantID, sessionID string, expectedGeneration uint64, resetRequestedAt *time.Time) (sessions.ManagedArchive, error) { - if err := s.checkExecutionAuthority(); err != nil { - return sessions.ManagedArchive{}, err - } - tenant, err := parseID(tenantID) - if err != nil { - return sessions.ManagedArchive{}, err - } - var result sessions.ManagedArchive - err = s.withPublicSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - // Session precedes deployment, matching Turn, allocation and input admission. - current, err := q.LockRuntimeDeployment(ctx) - if err != nil { - return err - } - if uint64(current.Generation) != expectedGeneration { - return &deployment.GenerationStaleError{CurrentGeneration: uint64(current.Generation)} - } - if !current.WebManaged || !current.InstallationID.Valid { - return deployment.ErrConflict - } - if current.ProviderKind == "" { - return deployment.ErrNotConfigured - } - environment, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: tenant, ID: session}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrInvalidInput - } - if err != nil { - return err - } - kind, err := sessions.EnvironmentType(environment.Configuration) - if err != nil || kind != "openai_hosted" { - return sessions.ErrInvalidInput - } - if resetRequestedAt != nil { - if !current.ResetClear.Valid || !current.ResetRequestedAt.Time.Equal(*resetRequestedAt) { - return deployment.ErrConflict - } - if current.ResetClear.String == "auto" { - busy, err := q.SessionBlocksAutoReset(ctx, session) - if err != nil { - return err - } - if busy { - return ErrSandboxResetSessionBusy - } - } - if environment.Environment.Status == "failed" || environment.Environment.Status == "expired" { - result, err = getManagedSessionArchive(ctx, q, tenant, session) - return err - } - source, err := deploymentpg.ResetSource(current) - if err != nil { - return err - } - project, err := q.GetSandboxResetProject(ctx, tenant) - if err != nil { - return err - } - source.ProjectID = runtimeUUID(project) - ctx = adminaudit.WithSource(ctx, source) - } - allocation, err := q.GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: tenant, EnvironmentID: environment.Environment.ID}) - allocated := err == nil - if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return err - } - if allocated && allocation.RuntimeAllocation.State != "released" && allocation.RuntimeAllocation.ProviderKey != current.InstallationID { - return deployment.ErrConflict - } - bound := sessionpg.BindSession(q, tenant, session) - if err := sessions.TrackInputActivity(ctx, bound, func(ctx context.Context) error { - if environment.Environment.Status != "failed" && environment.Environment.Status != "expired" { - if err := q.SetEnvironmentConnectionStatus(ctx, sqlc.SetEnvironmentConnectionStatusParams{ID: environment.Environment.ID, Status: "expired"}); err != nil { - return err - } - } - return sessions.CancelWork(ctx, bound) - }); err != nil { - return err - } - if allocated && allocation.RuntimeAllocation.State != "released" { - if _, err := q.RevokeArchivedRuntimeDevice(ctx, sqlc.RevokeArchivedRuntimeDeviceParams{TenantID: tenant, DeviceID: allocation.RuntimeAllocation.DeviceID, SessionID: session}); err != nil { - return err - } - if _, err := q.RequestRuntimeCleanup(ctx, allocation.RuntimeAllocation.ID); err != nil { - return err - } - } else if !allocated { - if err := q.ReleaseUnallocatedRuntimePlacement(ctx, session); err != nil { - return err - } - } - if err := auditpg.RecordAdminMutation(ctx, q, tenantID, "archive", "session", runtimeUUID(session)); err != nil { - return err - } - result, err = getManagedSessionArchive(ctx, q, tenant, session) - return err - }) - return result, err -} - -// GetManagedSessionArchive reads one database snapshot and never contacts compute. -func (s *Store) GetManagedSessionArchive(ctx context.Context, tenantID, sessionID string) (sessions.ManagedArchive, error) { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.ManagedArchive{}, err - } - return getManagedSessionArchive(ctx, s.queries, tenant, pgunit.PathID(sessionID)) -} - -func getManagedSessionArchive(ctx context.Context, q *sqlc.Queries, tenant, session pgtype.UUID) (sessions.ManagedArchive, error) { - row, err := q.GetManagedSessionArchive(ctx, sqlc.GetManagedSessionArchiveParams{TenantID: tenant, ID: session}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ManagedArchive{}, sessions.ErrNotFound - } - if err != nil { - return sessions.ManagedArchive{}, err - } - if row.EnvironmentType != "openai_hosted" { - return sessions.ManagedArchive{}, sessions.ErrInvalidInput - } - return sessions.ManagedArchive{SessionID: runtimeUUID(row.SessionID), EnvironmentID: runtimeUUID(row.EnvironmentID), State: row.State}, nil -} diff --git a/services/core/internal/store/admin_session_archive_race_test.go b/services/core/internal/store/admin_session_archive_race_test.go index 8a19c4d98..6903443af 100644 --- a/services/core/internal/store/admin_session_archive_race_test.go +++ b/services/core/internal/store/admin_session_archive_race_test.go @@ -36,7 +36,7 @@ func TestManagedSessionArchiveReleasesPendingNodePlacement(t *testing.T) { } onlineManagerNode(t, s, nodeID) tenant, session := managedArchiveSession(t, s, managerSessionInput(uuid.NewString())) - result, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1) + result, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1) if err != nil || result.State != "released" { t.Fatal(result, err) } @@ -67,7 +67,7 @@ func TestManagedSessionArchiveOrdersConcurrentInput(t *testing.T) { go func() { defer wg.Done() <-start - _, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1) + _, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1) if err != nil { t.Error(err) } @@ -78,7 +78,7 @@ func TestManagedSessionArchiveOrdersConcurrentInput(t *testing.T) { if err := s.pool.QueryRow(t.Context(), "SELECT count(*) FROM environment_input_reservations WHERE session_id=$1 AND state='pending'", session.ID).Scan(&pending); err != nil || pending != 0 { t.Fatal("input survived concurrent archive", pending, err) } - if row, err := s.GetSession(t.Context(), tenant, session.ID); err != nil || row.Environment.Status != "expired" { + if row, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID); err != nil || row.Environment.Status != "expired" { t.Fatal("concurrent input revived Environment", row, err) } } @@ -87,7 +87,7 @@ func TestManagedSessionArchiveOrdersConcurrentInput(t *testing.T) { func TestManagedSessionArchiveRejectsFileManagedDeployment(t *testing.T) { s, w, _ := managerFixture(t, 1, 1) tenant, session := managedArchiveSession(t, s, managerSessionInput(uuid.NewString())) - if _, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 0); !errors.Is(err, deployment.ErrConflict) { + if _, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 0); !errors.Is(err, deployment.ErrConflict) { t.Fatal("archive accepted file-managed deployment", err) } } diff --git a/services/core/internal/store/admin_session_archive_test.go b/services/core/internal/store/admin_session_archive_test.go index 511a6f458..d35a6477a 100644 --- a/services/core/internal/store/admin_session_archive_test.go +++ b/services/core/internal/store/admin_session_archive_test.go @@ -67,37 +67,35 @@ func TestManagedSessionArchiveUnallocatedAndGuards(t *testing.T) { input.InitialInputs = []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"waiting"}`)}} tenant, session := managedArchiveSession(t, s, input) ctx := adminDeleteContext(t.Context(), tenant, uuid.NewString()) - active, err := s.GetManagedSessionArchive(t.Context(), tenant, session.ID) + active, err := sessionAdapter(s).GetManagedSessionArchive(t.Context(), tenant, session.ID) if err != nil || active.State != "active" || active.SessionID != session.ID || active.EnvironmentID != session.Environment.ID { t.Fatal("unallocated Session status", active, err) } for _, generation := range []uint64{0, 2, ^uint64(0)} { - if _, err := w.ArchiveManagedSession(ctx, tenant, session.ID, generation); !errors.Is(err, deployment.ErrConflict) { + var stale *deployment.GenerationStaleError + if _, err := deploymentExecution(t, w).ArchiveSession(ctx, tenant, session.ID, generation); !errors.As(err, &stale) || stale.CurrentGeneration != 1 || !errors.Is(err, deployment.ErrConflict) { t.Fatal("archive accepted wrong generation", generation, err) } } - if _, err := s.ArchiveManagedSession(ctx, tenant, session.ID, 1); !errors.Is(err, ErrExecutionAuthority) { - t.Fatal("unleased archive accepted", err) - } for _, other := range []string{uuid.NewString(), "malformed"} { - if _, err := w.ArchiveManagedSession(ctx, tenant, other, 1); !errors.Is(err, sessions.ErrNotFound) { + if _, err := deploymentExecution(t, w).ArchiveSession(ctx, tenant, other, 1); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unknown archive", err) } - if _, err := s.GetManagedSessionArchive(ctx, tenant, other); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetManagedSessionArchive(ctx, tenant, other); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unknown status", err) } } - if _, err := w.ArchiveManagedSession(ctx, uuid.NewString(), session.ID, 1); !errors.Is(err, sessions.ErrNotFound) { + if _, err := deploymentExecution(t, w).ArchiveSession(ctx, uuid.NewString(), session.ID, 1); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign archive", err) } - if _, err := s.GetManagedSessionArchive(ctx, uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetManagedSessionArchive(ctx, uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign status", err) } - result, err := w.ArchiveManagedSession(ctx, tenant, session.ID, 1) + result, err := deploymentExecution(t, w).ArchiveSession(ctx, tenant, session.ID, 1) if err != nil || result.State != "released" { t.Fatal("unallocated archive", result, err) } - row, err := s.GetSession(t.Context(), tenant, session.ID) + row, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || row.Environment.Status != "expired" || row.EnvironmentFailure != nil || row.PendingInput || row.EnvironmentInputActivity != nil || row.LastTurn != nil { t.Fatal("archive fabricated failed execution", row, err) } @@ -106,11 +104,11 @@ func TestManagedSessionArchiveUnallocatedAndGuards(t *testing.T) { t.Fatal("archive left pending initial input", reservationState, err) } before := adminMutationSnapshot(t, s, "sessions", "environments", "environment_input_reservations", "session_events") - retry, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1) + retry, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1) if err != nil || retry != result || !reflect.DeepEqual(before, adminMutationSnapshot(t, s, "sessions", "environments", "environment_input_reservations", "session_events")) { t.Fatal("archive retry changed Session history", retry, err) } - if status, err := s.GetManagedSessionArchive(ctx, tenant, session.ID); err != nil || status != result { + if status, err := sessionAdapter(s).GetManagedSessionArchive(ctx, tenant, session.ID); err != nil || status != result { t.Fatal("status differs from committed archive", status, err) } if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, installation, runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, deployment.ErrInvalidInput) { @@ -146,7 +144,7 @@ func TestManagedSessionArchiveRetainsHistoryAndSettledResources(t *testing.T) { transition(t, w, tenant, session.ID, input.TurnID, sessions.TurnInProgress, sessions.TurnCompleted) history := adminMutationSnapshot(t, s, "sessions", "turns", "session_items", "session_artifacts", "source_files", "pg_largeobject", "pg_largeobject_metadata") request := uuid.NewString() - result, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, request), tenant, session.ID, 1) + result, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), tenant, request), tenant, session.ID, 1) if err != nil || result.State != "cleanup_pending" { t.Fatal(result, err) } @@ -167,7 +165,7 @@ func TestManagedSessionArchiveRetainsHistoryAndSettledResources(t *testing.T) { if _, err := deploymentExecution(t, w).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } - current, err := s.GetSession(t.Context(), tenant, session.ID) + current, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || current.EnvironmentFailure != nil || current.LastTurn == nil || current.LastTurn.Status != sessions.TurnCompleted || current.Environment.Status != "expired" { t.Fatal("cleanup rewrote completed outcome", current, err) } @@ -177,7 +175,7 @@ func TestManagedSessionArchiveRetainsHistoryAndSettledResources(t *testing.T) { if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } - if result, err := s.GetManagedSessionArchive(t.Context(), tenant, session.ID); err != nil || result.State != "released" { + if result, err := sessionAdapter(s).GetManagedSessionArchive(t.Context(), tenant, session.ID); err != nil || result.State != "released" { t.Fatal("release not reflected", result, err) } page, err := sessionAdapter(s).ListSessionArtifacts(t.Context(), tenant, session.ID, "", "", 100, true) @@ -214,12 +212,12 @@ func TestManagedSessionArchiveAuditFailureRollsBack(t *testing.T) { tables := []string{"sessions", "environments", "turns", "session_events", "devices", "runtime_allocations", "runtime_placements", "environment_input_reservations", "admin_audit_log"} before := adminMutationSnapshot(t, s, tables...) count := adminAuditRejections(t, s) - _, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, rejectedAdminRequest), tenant, session.ID, 1) + _, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), tenant, rejectedAdminRequest), tenant, session.ID, 1) requireAdminAuditFailure(t, s, err, count) if !reflect.DeepEqual(before, adminMutationSnapshot(t, s, tables...)) { t.Fatal("failed audit retained archive, revocation or cancellation") } - if _, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1); err != nil { + if _, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1); err != nil { t.Fatal(err) } turn, err := sessionAdapter(s).GetTurn(t.Context(), tenant, session.ID, input.TurnID) @@ -238,23 +236,23 @@ func TestManagedSessionArchivePreservesFailuresAndRejectsSelfHosted(t *testing.T if err := sessionExecution(t, w.lease).FailEnvironmentInitialization(t.Context(), sessions.EnvironmentInitialization{EnvironmentID: owner.EnvironmentID, SessionID: owner.SessionID, TenantID: owner.TenantID, DeviceID: owner.DeviceID}, sessions.ProvisioningFailure{Step: sessions.ProvisioningSetupCommand, Index: 0, ExitCode: 2}); err != nil { t.Fatal(err) } - failed, err := s.GetSession(t.Context(), tenant, session.ID) + failed, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || failed.EnvironmentFailure == nil { t.Fatal("failure fixture", err) } otherTenant, selfHosted := managedArchiveSession(t, s, environmentInput(uuid.NewString(), "self_hosted", "/workspace")) - if _, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1); err != nil { + if _, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1); err != nil { t.Fatal(err) } - after, err := s.GetSession(t.Context(), tenant, session.ID) + after, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || !reflect.DeepEqual(after.EnvironmentFailure, failed.EnvironmentFailure) || after.Environment.Status != "failed" { t.Fatal("archive changed recorded provisioning failure", after, err) } before := adminMutationSnapshot(t, s, "sessions", "environments", "admin_audit_log") - if _, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), otherTenant, uuid.NewString()), otherTenant, selfHosted.ID, 1); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), otherTenant, uuid.NewString()), otherTenant, selfHosted.ID, 1); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal("self-hosted archive accepted", err) } - if _, err := s.GetManagedSessionArchive(t.Context(), otherTenant, selfHosted.ID); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := sessionAdapter(s).GetManagedSessionArchive(t.Context(), otherTenant, selfHosted.ID); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal("self-hosted cleanup projected", err) } if !reflect.DeepEqual(before, adminMutationSnapshot(t, s, "sessions", "environments", "admin_audit_log")) { diff --git a/services/core/internal/store/admin_session_archive_worker_http_test.go b/services/core/internal/store/admin_session_archive_worker_http_test.go index df4b565e1..3c86b135c 100644 --- a/services/core/internal/store/admin_session_archive_worker_http_test.go +++ b/services/core/internal/store/admin_session_archive_worker_http_test.go @@ -11,6 +11,7 @@ import ( "sync" "testing" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/e2b" @@ -27,8 +28,9 @@ import ( "github.com/google/uuid" ) -// Use the original Store for administrator reads, as server startup does, and -// route archive writes through the real Worker's execution-owned Store clone. +// Serve administrator reads from the pooled Session adapter and archive through +// the deployment execution operations on the Worker's lease, as server startup +// does. func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { _, pool := store.NewManagedTestStore(t) cipher, err := credentialcrypto.New(bytes.Repeat([]byte{83}, 32)) @@ -52,7 +54,15 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { }, func(_ context.Context, setup deployment.Setup) (execution.PreparedRuntimeDeployment, error) { return execution.PreparedRuntimeDeployment{Config: providerConfig(setup)}, nil }) - worker := startWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}) + lease, err := pgunit.AcquireLease(t.Context(), pool) + if err != nil { + t.Fatal(err) + } + owner, err := fixtureOwner(db, s, lease) + if err != nil { + t.Fatal(errors.Join(err, lease.Close(t.Context()))) + } + worker := startOwnedWorker(t, t.Context(), db, &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), ManagedRuntimes: configuration}, owner) var once sync.Once stop := func() { once.Do(func() { @@ -82,7 +92,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { return session } active, untouched := create(), create() - owner, err := worker.ProvisionEnvironment(t.Context(), project.TenantID, active.Environment.ID, installation) + allocated, err := worker.ProvisionEnvironment(t.Context(), project.TenantID, active.Environment.ID, installation) if err != nil { t.Fatal(err) } @@ -90,14 +100,11 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.ArchiveManagedSession(ctx, project.TenantID, active.ID, 1); !errors.Is(err, store.ErrExecutionAuthority) { - t.Fatal("fixture admission Store unexpectedly holds execution ownership", err) - } admin, err := api.NewDeploymentAuthenticator([]string{runtimedevice.HashCredential("archive-administrator")}) if err != nil { t.Fatal(err) } - handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), workerExecution(worker), withCoreKeys(admin)) + handler, err := publicHandler(t, s, db, nil, "codex", storeKeys(s), workerExecution(t, worker), func(d *api.Dependencies) { d.Execution.SessionArchive = owner.Deployment }, withCoreKeys(admin)) if err != nil { t.Fatal(err) } @@ -113,10 +120,10 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { w := request(http.MethodPost, active.ID) var archived sessions.ManagedArchive if w.Code != http.StatusOK || json.Unmarshal(w.Body.Bytes(), &archived) != nil || archived.State != "cleanup_pending" || archived.SessionID != active.ID { - t.Fatalf("archive did not use Worker's leased Store: %d %s", w.Code, w.Body) + t.Fatalf("archive failed: %d %s", w.Code, w.Body) } allocation, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: project.TenantID, EnvironmentID: active.Environment.ID}) - if err != nil || allocation.ID != owner.ID || allocation.State != "cleanup_pending" { + if err != nil || allocation.ID != allocated.ID || allocation.State != "cleanup_pending" { t.Fatal("archive did not retain cleanup ownership", allocation, err) } turn, err := store.SessionAdapter(s).GetTurn(t.Context(), project.TenantID, active.ID, input.TurnID) diff --git a/services/core/internal/store/admin_summary.go b/services/core/internal/store/admin_summary.go index 97fffd512..d7578aaf0 100644 --- a/services/core/internal/store/admin_summary.go +++ b/services/core/internal/store/admin_summary.go @@ -51,7 +51,7 @@ func (s *Store) ReadAdminSummary(ctx context.Context, tenantID string, filter Ad if err != nil { return err } - session, err = readSessionActivity(ctx, q, session) + session, err = sessionpg.LoadSessionActivity(ctx, q, session) if err != nil { return err } diff --git a/services/core/internal/store/agent_execution_defaults_http_test.go b/services/core/internal/store/agent_execution_defaults_http_test.go index 95b9ead57..7b191ee3a 100644 --- a/services/core/internal/store/agent_execution_defaults_http_test.go +++ b/services/core/internal/store/agent_execution_defaults_http_test.go @@ -74,7 +74,7 @@ func TestAgentExecutionDefaultsPublicSnapshotAndPrecedence(t *testing.T) { sessionID := id(call("POST", "/v1/agents/sessions", body, key, 201)) assertSnapshot := func(sessionID, model, endpoint, key string) { t.Helper() - session, err := st.GetSession(t.Context(), tenant, sessionID) + session, err := store.SessionAdapter(st).GetSession(t.Context(), tenant, sessionID) if err != nil { t.Fatal(err) } @@ -105,8 +105,9 @@ func TestAgentExecutionDefaultsPublicSnapshotAndPrecedence(t *testing.T) { nativeHarness := strings.TrimSuffix(crossHarness, "}") + `,"x_agents_core":{"model_provider":{"protocol":"anthropic","base_url":"https://override.example/v1","api_key":"override-canary"}}}` created := id(call("POST", "/v1/agents/sessions", nativeHarness, uuid.NewString(), 201)) assertSnapshot(created, "model-override", "https://override.example/v1", "override-canary") - resolved, err := st.GetSession(t.Context(), tenant, created) - frozen, providerErr := store.SessionAdapter(st).SessionModelExecution(t.Context(), tenant, created) + reads := store.SessionAdapter(st) + resolved, err := reads.GetSession(t.Context(), tenant, created) + frozen, providerErr := reads.SessionModelExecution(t.Context(), tenant, created) if err != nil || providerErr != nil || resolved.Engine != "claude_sdk" || frozen == nil || frozen.Protocol != "anthropic" { t.Fatal("harness override did not freeze the native provider", err, providerErr) } diff --git a/services/core/internal/store/archive_cancellation_test.go b/services/core/internal/store/archive_cancellation_test.go index 4c173ad2b..9c56f82eb 100644 --- a/services/core/internal/store/archive_cancellation_test.go +++ b/services/core/internal/store/archive_cancellation_test.go @@ -39,7 +39,6 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { heartbeat := scenario != "receipt_without_heartbeat" s, db := newManagedTestStoreDB(t) leased := executionOwner(t, db, s) - writer := leased.Store t.Cleanup(func() { if err := leased.Lease.Close(context.Background()); err != nil { t.Error(err) @@ -158,7 +157,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { go func() { revokeDone <- fixtureSessionService(t, db).RevokeDevice(t.Context(), h.tenant, owner.DeviceID) }() } - archived, err := writer.ArchiveManagedSession(auditCtx, h.tenant, session.ID, 1) + archived, err := leased.Deployment.ArchiveSession(auditCtx, h.tenant, session.ID, 1) if err != nil || archived.State != "cleanup_pending" { t.Fatal(archived, err) } @@ -175,7 +174,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { // A repeat archive and cleanup must not recreate an explicitly // cleared marker, nor erase the marker from a fresh archive. repeatAudit := adminaudit.WithSource(t.Context(), adminaudit.Source{CredentialID: "fixture-admin", ProjectID: projectID, RequestID: uuid.NewString(), TraceID: uuid.NewString()}) - if _, err := writer.ArchiveManagedSession(repeatAudit, h.tenant, session.ID, 1); err != nil { + if _, err := leased.Deployment.ArchiveSession(repeatAudit, h.tenant, session.ID, 1); err != nil { t.Fatal(err) } if _, err := leased.Deployment.RequestCleanup(t.Context(), owner); err != nil { diff --git a/services/core/internal/store/archived_cancellation_migration_test.go b/services/core/internal/store/archived_cancellation_migration_test.go index bd966700b..c3f53cb0c 100644 --- a/services/core/internal/store/archived_cancellation_migration_test.go +++ b/services/core/internal/store/archived_cancellation_migration_test.go @@ -19,7 +19,7 @@ func TestArchivedCancellationMigrationDoesNotAdoptOldRevocations(t *testing.T) { input := submitMessage(t, s, tenant, session.ID, "waiting") transition(t, w, tenant, session.ID, input.TurnID, sessions.TurnQueued, sessions.TurnInProgress) transition(t, w, tenant, session.ID, input.TurnID, sessions.TurnInProgress, sessions.TurnWaiting) - if _, err := w.ArchiveManagedSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1); err != nil { + if _, err := deploymentExecution(t, w).ArchiveSession(adminDeleteContext(t.Context(), tenant, uuid.NewString()), tenant, session.ID, 1); err != nil { t.Fatal(err) } db := sql.OpenDB(stdlib.GetConnector(*s.pool.Config().ConnConfig)) diff --git a/services/core/internal/store/command_output_test.go b/services/core/internal/store/command_output_test.go index b8e90e586..e2e1cb94f 100644 --- a/services/core/internal/store/command_output_test.go +++ b/services/core/internal/store/command_output_test.go @@ -42,7 +42,7 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { t.Fatal(err) } } - before, _ := s.SessionEventCursor(ctx, tenant, session.ID) + before, _ := store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) // A bad command reference rolls back preceding valid fragments and their events. if err := journal.AppendTurnEvents(ctx, tenant, session.ID, input.TurnID, 4, []sessions.ExecutionEvent{ event("command_output", `{"id":"cmd","delta":"rollback"}`), @@ -50,7 +50,7 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { }); err == nil { t.Fatal("unknown command accepted") } - after, _ := s.SessionEventCursor(ctx, tenant, session.ID) + after, _ := store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if before != after { t.Fatal("rollback published output") } @@ -75,7 +75,7 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { } // Reopening the Store recovers committed Items without creating events. reopened := store.New(pool) - before, _ = s.SessionEventCursor(ctx, tenant, session.ID) + before, _ = store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) page, err = sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) if err != nil || len(page.Items) != 3 { t.Fatalf("recovery: %+v %v", page, err) @@ -83,11 +83,11 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { if page.Items[1].Status != "completed" || page.Items[1].Output != "authoritative" || page.Items[2].Status != "incomplete" || page.Items[2].Output != "已观察\n" { t.Fatal("completion/cancellation lost command output", page.Items) } - after, _ = s.SessionEventCursor(ctx, tenant, session.ID) + after, _ = store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if before != after { t.Fatal("query replayed events") } - if _, err := reopened.ListSessionEvents(ctx, uuid.NewString(), session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { + if _, err := store.SessionAdapter(reopened).ListSessionEvents(ctx, uuid.NewString(), session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign event access", err) } var fragments []string @@ -95,7 +95,7 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { indexes := map[string]int32{} cursor := int64(0) for { - changes, err := reopened.ListSessionEvents(ctx, tenant, session.ID, cursor) + changes, err := store.SessionAdapter(reopened).ListSessionEvents(ctx, tenant, session.ID, cursor) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/creation_stream_settlement_public_test.go b/services/core/internal/store/creation_stream_settlement_public_test.go index e50fd5360..7de1cccaf 100644 --- a/services/core/internal/store/creation_stream_settlement_public_test.go +++ b/services/core/internal/store/creation_stream_settlement_public_test.go @@ -13,6 +13,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -173,7 +174,7 @@ func TestCreationStreamPublicLifetimes(t *testing.T) { if _, err := s.ReserveEnvironmentInput(t.Context(), tenant, first.Session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}); err != nil { t.Fatal(err) } - if current, err := s.GetSession(t.Context(), tenant, first.Session.ID); err != nil || !current.PendingInput { + if current, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, first.Session.ID); err != nil || !current.PendingInput { t.Fatal("later input is not pending", err) } retry := openStream(t, server, token, http.MethodPost, "/v1/agents/sessions", idle, "no-input") @@ -206,14 +207,14 @@ func TestCreationStreamPublicLifetimes(t *testing.T) { if err := db.pool.QueryRow(t.Context(), "SELECT id FROM environment_input_reservations WHERE session_id=$1 AND is_initial", session).Scan(&reservation); err != nil { t.Fatal(err) } - cursor, err := s.SessionEventCursor(t.Context(), tenant, session) + cursor, err := store.SessionAdapter(s).SessionEventCursor(t.Context(), tenant, session) if err != nil { t.Fatal(err) } if settled, err := s.CancelEnvironmentInput(t.Context(), tenant, session, reservation); err != nil || settled.State != sessions.EnvironmentInputCancelled { t.Fatal(settled.State, err) } - if after, err := s.SessionEventCursor(t.Context(), tenant, session); err != nil || after != cursor { + if after, err := store.SessionAdapter(s).SessionEventCursor(t.Context(), tenant, session); err != nil || after != cursor { t.Fatal("cancellation recorded a Session event; this case needs a silent settlement", after, cursor, err) } fresh.ended(t, 5*time.Second) diff --git a/services/core/internal/store/deployment_model_providers_http_test.go b/services/core/internal/store/deployment_model_providers_http_test.go index 63866bca3..0b0848aa9 100644 --- a/services/core/internal/store/deployment_model_providers_http_test.go +++ b/services/core/internal/store/deployment_model_providers_http_test.go @@ -194,7 +194,7 @@ func TestDeploymentModelProvidersHTTP(t *testing.T) { if providerOf(hostedID) != "deployment-canary" { t.Fatal("hosted Session did not freeze the deployment default") } - projection, err := st.GetSessionExecutionConfiguration(t.Context(), tenant, hostedID) + projection, err := store.SessionAdapter(st).GetSessionExecutionConfiguration(t.Context(), tenant, hostedID) if err != nil || projection.ModelProvider.Source != "deployment" || projection.ModelProvider.Status != "available" || projection.ModelProvider.Configuration == nil || projection.ModelProvider.Configuration.BaseURL != "https://deployment.example/v1" { t.Fatal("deployment selection not recorded", projection, err) } @@ -271,7 +271,7 @@ func TestLegacySessionWithoutProviderCannotStartWork(t *testing.T) { got, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, legacy.ID, pending.ID) return err == nil && got.State == sessions.EnvironmentInputFailed }) - session, err := h.s.GetSession(t.Context(), h.tenant, legacy.ID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, legacy.ID) if err != nil || session.EnvironmentInputActivity == nil || session.EnvironmentInputActivity.Failure != "model_provider_required" { t.Fatal("legacy reservation did not fail with its reason", session.EnvironmentInputActivity, err) } diff --git a/services/core/internal/store/environment_admission_test.go b/services/core/internal/store/environment_admission_test.go index 5aae0496f..48f5a5d89 100644 --- a/services/core/internal/store/environment_admission_test.go +++ b/services/core/internal/store/environment_admission_test.go @@ -92,7 +92,7 @@ func TestEnvironmentAdmissionWaitsForPreparedClaimAndRetainsRetry(t *testing.T) default: } } - session, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil || session.LastTurn != nil || session.EnvironmentInputActivity == nil || session.EnvironmentInputActivity.Status != "requires_action" { t.Fatal("waiting activity", session, err) } @@ -186,7 +186,7 @@ func TestEnvironmentAdmissionSettlementDoesNotCreateTurn(t *testing.T) { } case "deleted": expected = sessions.ErrNotFound - if err := h.s.DeleteSession(t.Context(), h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := store.SessionService(t, h.s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("pending input deleted", err) } if err := h.s.CommitLegacyDeletion(t.Context(), h.tenant, h.session.ID); err != nil { diff --git a/services/core/internal/store/environment_claim_worker_test.go b/services/core/internal/store/environment_claim_worker_test.go index 8efa54b68..a3b57b5b2 100644 --- a/services/core/internal/store/environment_claim_worker_test.go +++ b/services/core/internal/store/environment_claim_worker_test.go @@ -26,7 +26,7 @@ func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { } turnID := got.Receipts[0].TurnID if deleted { - if err := s.DeleteSession(t.Context(), tenant, pending.SessionID); !errors.Is(err, sessions.ErrNotIdle) { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: pending.SessionID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("claimed Session deleted", err) } if err := s.CommitLegacyDeletion(t.Context(), tenant, pending.SessionID); err != nil { diff --git a/services/core/internal/store/environment_connection_events_test.go b/services/core/internal/store/environment_connection_events_test.go index 135b9c5a7..f5dfc950f 100644 --- a/services/core/internal/store/environment_connection_events_test.go +++ b/services/core/internal/store/environment_connection_events_test.go @@ -39,7 +39,7 @@ func retainedEnvironmentEvents(t *testing.T, ctx context.Context, s *store.Store var after int64 events := []v1.SessionEvent{} for { - changes, err := s.ListSessionEvents(ctx, tenant, session, after) + changes, err := store.SessionAdapter(s).ListSessionEvents(ctx, tenant, session, after) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_connections_test.go b/services/core/internal/store/environment_connections_test.go index 5c8fe9407..79242db7c 100644 --- a/services/core/internal/store/environment_connections_test.go +++ b/services/core/internal/store/environment_connections_test.go @@ -37,7 +37,7 @@ func connectionSnapshot(t *testing.T, pool *pgxpool.Pool, id string) string { func connectionChanges(t *testing.T, s *Store, tenant, session string) []sessions.SessionChange { t.Helper() - all, err := s.ListSessionEvents(t.Context(), tenant, session, 0) + all, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session, 0) if err != nil { t.Fatal(err) } @@ -186,7 +186,7 @@ func TestEnvironmentConnectionDoesNotReviveDeletedOrTerminalResources(t *testing t.Fatal(err) } if status == "deleted" { - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } } else { diff --git a/services/core/internal/store/environment_device_test.go b/services/core/internal/store/environment_device_test.go index c9578f93d..46801cb71 100644 --- a/services/core/internal/store/environment_device_test.go +++ b/services/core/internal/store/environment_device_test.go @@ -37,7 +37,7 @@ func TestWorkerEnvironmentSelectsCapableDeviceWithoutMovingBinding(t *testing.T) if _, err := fixtureSessionStore(h.db).GetSessionDevice(t.Context(), h.tenant, pending.SessionID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("unregistered Runtime was assigned general compute", err) } - session, err := h.s.GetSession(t.Context(), h.tenant, pending.SessionID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, pending.SessionID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_directory_test.go b/services/core/internal/store/environment_directory_test.go index d98cf67a4..4c24fa01d 100644 --- a/services/core/internal/store/environment_directory_test.go +++ b/services/core/internal/store/environment_directory_test.go @@ -130,7 +130,7 @@ func TestEnvironmentDirectoryWorkerReadsWithoutExecutionPrerequisites(t *testing if got.err != nil || len(got.value.Entries) != 1 { t.Fatal("directory result", got.err) } - session, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil || session.LastTurn != nil || session.EnvironmentInputActivity != nil { t.Fatal("directory read manufactured execution") } diff --git a/services/core/internal/store/environment_executor_credentials_test.go b/services/core/internal/store/environment_executor_credentials_test.go index 033ec501f..ad1de8b60 100644 --- a/services/core/internal/store/environment_executor_credentials_test.go +++ b/services/core/internal/store/environment_executor_credentials_test.go @@ -110,7 +110,7 @@ func TestEnvironmentExecutorCredentialLifecycle(t *testing.T) { } check(restarted, next.Token, false) check(restarted, restored.Token, true) - if err := s.DeleteSession(ctx, tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } check(restarted, restored.Token, false) @@ -172,7 +172,7 @@ func TestEnvironmentExecutorConcurrentIssueAndDeletion(t *testing.T) { t.Error(err) } }() - if err := s.DeleteSession(ctx, tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } wg.Wait() diff --git a/services/core/internal/store/environment_executor_management_test.go b/services/core/internal/store/environment_executor_management_test.go index 3e0b373d8..e9f3bf138 100644 --- a/services/core/internal/store/environment_executor_management_test.go +++ b/services/core/internal/store/environment_executor_management_test.go @@ -135,7 +135,7 @@ func TestProjectEnvironmentExecutorManagement(t *testing.T) { t.Fatal("audit actions", actions, rows.Err()) } - if err := s.DeleteSession(ctx, p.TenantID, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: p.TenantID, SessionID: session.ID}); err != nil { t.Fatal(err) } if _, err := projectCredentials(ctx, s, p, one.ID); !errors.Is(err, sessions.ErrNotFound) { 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 ef032e80e..7076eaa71 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 @@ -60,7 +60,7 @@ func TestEnvironmentFileCreateRejectionsLeaveNoReceiptOrConsumption(t *testing.T {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "tenant-b", TokenSHA256: runtimedevice.HashCredential(other), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(w)) + handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(t, w)) if err != nil { t.Fatal(err) } @@ -169,7 +169,7 @@ func TestEnvironmentFileCreateRejectionsLeaveNoReceiptOrConsumption(t *testing.T if got := states(); !reflect.DeepEqual(got, map[string]int{"rejected": 3, "committed": 1}) { t.Fatal("successor receipt", got) } - session, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil || session.LastTurn != nil { t.Fatal("file writes created model execution", err) } diff --git a/services/core/internal/store/environment_file_writes_test.go b/services/core/internal/store/environment_file_writes_test.go index d4fb1b3a0..68cc9c7c6 100644 --- a/services/core/internal/store/environment_file_writes_test.go +++ b/services/core/internal/store/environment_file_writes_test.go @@ -70,7 +70,7 @@ func TestEnvironmentFileWriteRetainsUnknownAcrossLeaseLoss(t *testing.T) { if _, err := sessionAdapter(reopened).GetEnvironment(ctx, f.tenant, f.env.ID); err != nil { t.Fatal("write gate prevented metadata read", err) } - if _, err := reopened.GetSession(ctx, f.tenant, f.session.ID); err != nil { + if _, err := sessionAdapter(reopened).GetSession(ctx, f.tenant, f.session.ID); err != nil { t.Fatal("write gate prevented recovery read", err) } if receipt, err := reopened.RequestCancel(ctx, f.tenant, f.session.ID, "idle-cancel"); err != nil || receipt.TurnID != "" { @@ -113,7 +113,7 @@ func TestEnvironmentFileWriteMatchesReceiptAndRetainsDeletedOwner(t *testing.T) t.Fatal("non-receipt settled write", state, err) } } - if err := f.s.DeleteSession(ctx, f.tenant, f.session.ID); err != nil { + if err := sessionService(t, f.s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: f.tenant, SessionID: f.session.ID}); err != nil { t.Fatal(err) } if got, err := FixtureFileWrite(ctx, f.s.pool, f.tenant, f.env.ID, f.key.ID); err != nil || got.State != "pending" { diff --git a/services/core/internal/store/environment_initial_input_test.go b/services/core/internal/store/environment_initial_input_test.go index 1a9b2badf..ae58bba5b 100644 --- a/services/core/internal/store/environment_initial_input_test.go +++ b/services/core/internal/store/environment_initial_input_test.go @@ -74,7 +74,7 @@ func TestEnvironmentInitialExpiryRollsBackWithFailureEventAndSerializesPromotion t.Fatal("expiry/promotion race started work", got) } } - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + events, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil || len(events) != 2 || events[1].Event.Type != "agent.session.failed" { t.Fatal("racing settlement duplicated or lost failure", events, err) } @@ -116,7 +116,7 @@ func TestEnvironmentInitialInputCreationRetainsCursorIdentityAndPromotion(t *tes status, actionEnvironment = "", "" } requireEnvironmentInputActivity(t, s, tenant, session.ID, status, actionEnvironment) - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, creation.Cursor) + events, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, creation.Cursor) expectedEvents := 1 if kind == "openai_hosted" { expectedEvents = 0 @@ -172,7 +172,7 @@ func TestEnvironmentInitialInputCreationRetainsCursorIdentityAndPromotion(t *tes t.Fatal(err) } environmentInputHistory(t, pool, session.ID, 1, 2) - after, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + after, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil || !reflect.DeepEqual(events, after[:len(events)]) { t.Fatal("initial snapshot changed after promotion", err) } @@ -221,7 +221,7 @@ func TestEnvironmentInitialInputExpiryHasNoTurnAndCannotReplay(t *testing.T) { if failed.Environment.Status != "pending" || failed.LastTurn != nil || failed.PendingInput || !failed.EnvironmentInputActivity.LastActiveAt.Equal(*reservation.SettledAt) { t.Fatal("input expiry changed Environment/Turn", failed) } - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + events, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) expectedEvents := 2 if kind == "openai_hosted" { expectedEvents = 1 @@ -261,21 +261,21 @@ func TestEnvironmentInitialInputExpiryHasNoTurnAndCannotReplay(t *testing.T) { t.Fatal("later submission inferred initial origin") } requireEnvironmentInputActivity(t, reopened, tenant, session.ID, "idle", "") - after, err := reopened.ListSessionEvents(t.Context(), tenant, session.ID, 0) + after, err := sessionAdapter(reopened).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil || len(after) < len(events) || !reflect.DeepEqual(events, after[:len(events)]) { t.Fatal("later work changed historical failure", err) } // The later input is still pending, so deletion waits for it to settle. - if err := reopened.DeleteSession(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := sessionService(t, reopened).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("pending later input deleted", err) } if _, err := reopened.CancelEnvironmentInput(t.Context(), tenant, session.ID, later.ID); err != nil { t.Fatal(err) } - if err := reopened.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := sessionService(t, reopened).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } - if _, err := reopened.GetSession(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(reopened).GetSession(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted initial Session remained visible", err) } }) diff --git a/services/core/internal/store/environment_initialization_test.go b/services/core/internal/store/environment_initialization_test.go index 05b4be525..874ef9507 100644 --- a/services/core/internal/store/environment_initialization_test.go +++ b/services/core/internal/store/environment_initialization_test.go @@ -160,7 +160,7 @@ func TestUserManagedPreparationUsesAuthenticatedRuntimeWithoutAllocation(t *test if peer.writes.Load() != expected { t.Fatal("failed, unavailable or revoked effect replayed", peer.writes.Load()) } - value, err := s.GetSession(t.Context(), principal.TenantID, session.ID) + value, err := store.SessionAdapter(s).GetSession(t.Context(), principal.TenantID, session.ID) if err != nil || value.EnvironmentFailure == nil { t.Fatal("failure lost shared Session state", err) } diff --git a/services/core/internal/store/environment_input_activity_test.go b/services/core/internal/store/environment_input_activity_test.go index af7f32f8d..2dd6a8cd9 100644 --- a/services/core/internal/store/environment_input_activity_test.go +++ b/services/core/internal/store/environment_input_activity_test.go @@ -14,7 +14,7 @@ import ( func requireEnvironmentInputActivity(t *testing.T, s *Store, tenant, session, status, environment string) sessions.Session { t.Helper() - value, err := s.GetSession(t.Context(), tenant, session) + value, err := sessionAdapter(s).GetSession(t.Context(), tenant, session) if err != nil { t.Fatal(err) } @@ -26,7 +26,7 @@ func requireEnvironmentInputActivity(t *testing.T, s *Store, tenant, session, st } else if activity == nil || activity.Status != status || activity.EnvironmentID != environment || activity.LastActiveAt.IsZero() { t.Fatal("input activity", activity, status, environment) } - page, err := s.ListSessions(t.Context(), tenant, "", 100, false, nil) + page, err := sessionAdapter(s).ListSessions(t.Context(), tenant, "", 100, false, nil) if err != nil || len(page.Sessions) != 1 || !reflect.DeepEqual(page.Sessions[0].EnvironmentInputActivity, activity) { t.Fatal("list and retrieve activity differ", err) } @@ -47,7 +47,7 @@ func TestEnvironmentInputActivityWaitsBeforeTurnAndClearsOnConnection(t *testing t.Fatal("waiting input fabricated a Turn") } environmentInputHistory(t, pool, session.ID, 0, 0) - first, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + first, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil || len(first) != 1 || first[0].Event.Type != "agent.session.requires_action" || first[0].Turn != nil || first[0].EnvironmentInputActivity == nil { t.Fatal("missing pre-Turn snapshot", first, err) } @@ -64,7 +64,7 @@ func TestEnvironmentInputActivityWaitsBeforeTurnAndClearsOnConnection(t *testing if idle.LastTurn != nil { t.Fatal("connection fabricated readiness or Turn") } - changes, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + changes, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil || len(changes) != 3 || changes[1].Event.Type != "agent.session.environment.connected" || changes[2].Event.Type != "agent.session.idle" { t.Fatal("connection/action order", changes, err) } @@ -92,7 +92,7 @@ func TestEnvironmentInputActivityWaitsBeforeTurnAndClearsOnConnection(t *testing if active.LastTurn == nil || active.LastTurn.Status != sessions.TurnInProgress { t.Fatal("normal Turn did not take ownership") } - if _, err := s.GetSession(t.Context(), uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSession(t.Context(), uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign Session activity visible", err) } } @@ -111,11 +111,11 @@ func TestEnvironmentInputActivitySettlementAndNewerWork(t *testing.T) { t.Fatal(err) } reservation := reserveEnvironmentInput(t, s, tenant, session.ID, "waiting") - value, err := s.GetSession(t.Context(), tenant, session.ID) + value, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || value.LastTurn.Status != sessions.TurnFailed || value.EnvironmentInputActivity.Status != "requires_action" || !value.PendingInput { t.Fatal("prior failure hid waiting input", err) } - waitingCursor, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + waitingCursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -143,16 +143,16 @@ func TestEnvironmentInputActivitySettlementAndNewerWork(t *testing.T) { // The settled later reservation no longer counts as pending input, so // creation streams can end instead of waiting for work that cannot start. settled := requireEnvironmentInputActivity(t, s, tenant, session.ID, "idle", "") - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, waitingCursor) + events, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, waitingCursor) if err != nil || settled.PendingInput || len(events) != 1 || events[0].Event.Type != "agent.session.idle" || events[0].Turn != nil || !events[0].Settled { t.Fatal("settled reservation remained pending", settled.EnvironmentInputActivity, events, err) } - cursor, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } reserveEnvironmentInput(t, s, tenant, session.ID, "waiting") - if after, err := s.SessionEventCursor(t.Context(), tenant, session.ID); err != nil || after != cursor { + if after, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID); err != nil || after != cursor { t.Fatal("settled retry repeated activity", after, cursor, err) } if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "newer", []sessions.Input{messageInput("newer")}); err != nil { @@ -231,27 +231,27 @@ func TestEnvironmentInputActivityRecoversWaitingActionAndHidesDeletion(t *testin if err != nil || got.State != sessions.EnvironmentInputPending || !got.Deadline.Equal(reservation.Deadline) { t.Fatal("recovery changed waiting input or its deadline", got, err) } - cursor, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } if err := sessionExecution(t, next.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 2, true); err != nil { t.Fatal(err) } - if after, err := s.SessionEventCursor(t.Context(), tenant, session.ID); err != nil || after != cursor { + if after, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID); err != nil || after != cursor { t.Fatal("retired generation changed activity", after, cursor, err) } environmentInputHistory(t, pool, session.ID, 0, 0) - if err := s.DeleteSession(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("waiting input deleted", err) } if err := s.commitLegacyDeletion(t.Context(), tenant, session.ID); err != nil { t.Fatal(err) } - if _, err := s.GetSession(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted activity remained visible", err) } - if _, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted activity events remained visible", err) } } diff --git a/services/core/internal/store/environment_input_claim_test.go b/services/core/internal/store/environment_input_claim_test.go index 4f588dad2..8ce51ee8d 100644 --- a/services/core/internal/store/environment_input_claim_test.go +++ b/services/core/internal/store/environment_input_claim_test.go @@ -12,7 +12,7 @@ func TestEnvironmentInputConcurrentPromotionClaimsOnce(t *testing.T) { writer := executionWriter(t, s) tenant, session := environmentInputSession(t, s) pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - reservationCursor, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + reservationCursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -58,7 +58,7 @@ func TestEnvironmentInputConcurrentPromotionClaimsOnce(t *testing.T) { t.Fatal("promotion did not persist its execution claim", turn, err) } environmentInputHistory(t, pool, session.ID, 1, 2) - changes, err := s.ListSessionEvents(t.Context(), tenant, session.ID, reservationCursor) + changes, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, reservationCursor) if err != nil || len(changes) < 2 { t.Fatal("missing promotion events", changes, err) } @@ -77,7 +77,7 @@ func TestEnvironmentInputConcurrentPromotionClaimsOnce(t *testing.T) { } transition(t, writer, tenant, session.ID, turnID, sessions.TurnInProgress, sessions.TurnCompleted) later := reserveEnvironmentInput(t, s, tenant, session.ID, "later") - cursor, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -85,7 +85,7 @@ func TestEnvironmentInputConcurrentPromotionClaimsOnce(t *testing.T) { if err != nil || len(retry.Receipts) != 2 || !retry.Receipts[0].Replayed || retry.Receipts[0].TurnID != turnID { t.Fatal("terminal retry reclaimed execution", retry, err) } - after, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + after, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil || after != cursor { t.Fatal("terminal retry published events", after, cursor, err) } diff --git a/services/core/internal/store/environment_input_expiry_test.go b/services/core/internal/store/environment_input_expiry_test.go index bcc6ecd97..f733964f8 100644 --- a/services/core/internal/store/environment_input_expiry_test.go +++ b/services/core/internal/store/environment_input_expiry_test.go @@ -112,7 +112,7 @@ func TestEnvironmentExpirySerializesWithTargetedSettlement(t *testing.T) { case "cancel": _, err = s.CancelEnvironmentInput(t.Context(), tenant, session.ID, pending.ID) case "delete": - err = s.DeleteSession(t.Context(), tenant, session.ID) + err = sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}) } results <- err }() diff --git a/services/core/internal/store/environment_input_settlement_test.go b/services/core/internal/store/environment_input_settlement_test.go index 82039e2d7..874d8958a 100644 --- a/services/core/internal/store/environment_input_settlement_test.go +++ b/services/core/internal/store/environment_input_settlement_test.go @@ -236,7 +236,7 @@ func TestEnvironmentInputDeletionSettlesPendingAndFencesPromotion(t *testing.T) }() } if !concurrent { - if err := s.DeleteSession(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("pending input deleted", err) } } diff --git a/services/core/internal/store/environment_installation_test.go b/services/core/internal/store/environment_installation_test.go index 20c4726e4..2b25b5e38 100644 --- a/services/core/internal/store/environment_installation_test.go +++ b/services/core/internal/store/environment_installation_test.go @@ -94,7 +94,7 @@ func TestEnvironmentInstallationClaimLifetimeAndRetries(t *testing.T) { 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 { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: p.TenantID, SessionID: session.ID}); err != nil { t.Fatal(err) } if _, err := installations.ValidateEnvironmentInstallation(ctx, token, "build"); !errors.Is(err, sessions.ErrInstallationAuthorization) { diff --git a/services/core/internal/store/environment_mcp_public_test.go b/services/core/internal/store/environment_mcp_public_test.go index 35a0c47d4..20e61f076 100644 --- a/services/core/internal/store/environment_mcp_public_test.go +++ b/services/core/internal/store/environment_mcp_public_test.go @@ -17,7 +17,7 @@ func TestPublicEnvironmentMCPUsesAttachedVaultSelection(t *testing.T) { t.Run(kind, func(t *testing.T) { s, db, tenant, vault, credential := selfHostedMCPAdmissionFixture(t) auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "test", TenantID: tenant, TokenSHA256: runtimedevice.HashCredential("test-token")}}) - handler, err := publicHandler(t, s, db, auth, kind, workerExecution(&execution.Worker{}), executorURL("https://executor.example")) + handler, err := publicHandler(t, s, db, auth, kind, workerExecution(t, &execution.Worker{}), executorURL("https://executor.example")) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_skill_references_test.go b/services/core/internal/store/environment_skill_references_test.go index a0820cf48..ce84a78b8 100644 --- a/services/core/internal/store/environment_skill_references_test.go +++ b/services/core/internal/store/environment_skill_references_test.go @@ -69,7 +69,7 @@ func TestSkillReferencesFreezeWithinSessionCreation(t *testing.T) { if err != nil || len(setup.Skills) != 1 || setup.Skills[0].Metadata.Type != "skill_reference" || setup.Skills[0].Metadata.SkillID != skill.ID || setup.Skills[0].Metadata.Version != version || !bytes.Equal(setup.Skills[0].Archive, archive) { t.Fatal("incorrect frozen installation", err) } - session, err := s.GetSession(t.Context(), tenant, id) + session, err := sessionAdapter(s).GetSession(t.Context(), tenant, id) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_work_test.go b/services/core/internal/store/environment_work_test.go index 2ba9a501c..5085d3585 100644 --- a/services/core/internal/store/environment_work_test.go +++ b/services/core/internal/store/environment_work_test.go @@ -27,7 +27,7 @@ func TestEnvironmentInputWorkFiltersAndPagesDevices(t *testing.T) { t.Fatal(err) } case "deleted": - if err := h.s.DeleteSession(t.Context(), h.tenant, pending.SessionID); !errors.Is(err, sessions.ErrNotIdle) { + if err := store.SessionService(t, h.s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: pending.SessionID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("pending input deleted", err) } if err := h.s.CommitLegacyDeletion(t.Context(), h.tenant, pending.SessionID); err != nil { diff --git a/services/core/internal/store/environment_worker_helpers_test.go b/services/core/internal/store/environment_worker_helpers_test.go index 8e6dcfefd..d39561867 100644 --- a/services/core/internal/store/environment_worker_helpers_test.go +++ b/services/core/internal/store/environment_worker_helpers_test.go @@ -34,7 +34,7 @@ func workerEnvironmentCapabilities() proto.AgentKindCapabilities { func workerEnvironmentReservation(t *testing.T, h *dispatchHarness) sessions.EnvironmentInputReservation { t.Helper() pending := unboundWorkerEnvironmentReservation(t, h) - session, err := h.s.GetSession(t.Context(), h.tenant, pending.SessionID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, pending.SessionID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_worker_scan_test.go b/services/core/internal/store/environment_worker_scan_test.go index 64fe1191f..f160dbb04 100644 --- a/services/core/internal/store/environment_worker_scan_test.go +++ b/services/core/internal/store/environment_worker_scan_test.go @@ -6,6 +6,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestWorkerEnvironmentRetriesNewlyReadyAtNextScan(t *testing.T) { @@ -59,7 +60,7 @@ func TestWorkerEnvironmentPaginationReachesReadyTail(t *testing.T) { last = pending } } - session, err := h.s.GetSession(t.Context(), h.tenant, last.SessionID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, last.SessionID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environments_test.go b/services/core/internal/store/environments_test.go index d423b4156..362c87b58 100644 --- a/services/core/internal/store/environments_test.go +++ b/services/core/internal/store/environments_test.go @@ -62,7 +62,7 @@ func TestEnvironmentOwnershipPersistsAndStaysScoped(t *testing.T) { if _, err := s.CreateSession(ctx, tenant, changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal("changed configuration accepted", err) } - if _, err := s.UpdateSessionMetadata(ctx, tenant, session.ID, map[string]string{"updated": "yes"}); err != nil { + if _, err := sessionService(t, s).UpdateSessionMetadata(ctx, sessions.UpdateSessionMetadataCommand{TenantID: tenant, SessionID: session.ID, Metadata: map[string]string{"updated": "yes"}}); err != nil { t.Fatal(err) } pool.Close() @@ -153,7 +153,7 @@ func TestEnvironmentCreationWinnerOwnsSnapshotAndIdentity(t *testing.T) { t.Fatal(reservations, err) } environmentInputHistory(t, pool, first.creation.Session.ID, 0, 0) - events, err := s.ListSessionEvents(ctx, tenant, first.creation.Session.ID, 0) + events, err := sessionAdapter(s).ListSessionEvents(ctx, tenant, first.creation.Session.ID, 0) if err != nil { t.Fatal(err) } @@ -174,7 +174,7 @@ func TestEnvironmentCreationWinnerOwnsSnapshotAndIdentity(t *testing.T) { if err != nil || !reflect.DeepEqual(environment, first.environment) { t.Fatal(environment, err) } - after, err := restarted.ListSessionEvents(ctx, tenant, found.Session.ID, 0) + after, err := sessionAdapter(restarted).ListSessionEvents(ctx, tenant, found.Session.ID, 0) if err != nil || !reflect.DeepEqual(events, after) { t.Fatal("retry emitted work", err) } @@ -242,7 +242,7 @@ func TestEnvironmentDeletionHidesWithoutDestroyingOwnership(t *testing.T) { if err != nil { t.Fatal(err) } - if err := s.DeleteSession(ctx, tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } if _, err := sessionAdapter(s).GetEnvironment(ctx, tenant, environment.ID); !errors.Is(err, sessions.ErrNotFound) { diff --git a/services/core/internal/store/execution_test.go b/services/core/internal/store/execution_test.go index f2675a51b..c50428880 100644 --- a/services/core/internal/store/execution_test.go +++ b/services/core/internal/store/execution_test.go @@ -50,7 +50,7 @@ func sessionExecution(t *testing.T, lease *pgunit.Lease) *sessions.ExecutionOper // pool. func transitionTurn(ctx context.Context, s *Store, tenant, session, turn string, transition sessions.TurnTransition) (sessions.Turn, error) { var moved sessions.Turn - err := s.withLockedSession(ctx, tenant, session, false, func(ctx context.Context, q *sqlc.Queries, id pgtype.UUID, _ sessions.LockedSession) error { + err := s.withSession(ctx, tenant, session, func(ctx context.Context, q *sqlc.Queries, id pgtype.UUID) error { owner, err := parseID(tenant) if err != nil { return err @@ -139,11 +139,11 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { if err = SubmitFixtureFunctionResult(t.Context(), s, tenant, waiting.ID, waitInput.TurnID, call.CallID, json.RawMessage(`{"success":true,"output":"saved"}`)); err != nil { t.Fatal(err) } - before, err := s.GetSession(t.Context(), tenant, active.ID) + before, err := sessionAdapter(s).GetSession(t.Context(), tenant, active.ID) if err != nil { t.Fatal(err) } - cursor, err := s.SessionEventCursor(t.Context(), tenant, active.ID) + cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, active.ID) if err != nil { t.Fatal(err) } @@ -172,11 +172,11 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { _, err = writer.ExpireEnvironmentInputs(t.Context()) mustReject("input expiry", err) mustReject("ownership check", writer.lease.CheckOwnership(t.Context())) - after, err := s.GetSession(t.Context(), tenant, active.ID) + after, err := sessionAdapter(s).GetSession(t.Context(), tenant, active.ID) if err != nil || !reflect.DeepEqual(before, after) { t.Fatal("stale state persisted", after, err) } - afterCursor, err := s.SessionEventCursor(t.Context(), tenant, active.ID) + afterCursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, active.ID) if err != nil || afterCursor != cursor { t.Fatal("stale events published", afterCursor, err) } @@ -290,29 +290,22 @@ func TestPooledStoreHasNoExecutionAuthority(t *testing.T) { s, _ := testStore(t) tenant, session := newTurnSession(t, s) submitMessage(t, s, tenant, session.ID, "start") - before, err := s.GetSession(t.Context(), tenant, session.ID) + before, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } - cursor, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } - _, archiveErr := s.ArchiveManagedSession(t.Context(), tenant, session.ID, 0) - _, expiryErr := s.ExpireEnvironmentInputs(t.Context()) - for name, err := range map[string]error{ - "archive": archiveErr, - "input expiry": expiryErr, - } { - if !errors.Is(err, ErrExecutionAuthority) { - t.Fatalf("pooled Store ran %s: %v", name, err) - } + if _, err := s.ExpireEnvironmentInputs(t.Context()); !errors.Is(err, ErrExecutionAuthority) { + t.Fatalf("pooled Store ran input expiry: %v", err) } - after, err := s.GetSession(t.Context(), tenant, session.ID) + after, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || !reflect.DeepEqual(before, after) { t.Fatal("rejected execution operation changed the Session", after, err) } - if next, err := s.SessionEventCursor(t.Context(), tenant, session.ID); err != nil || next != cursor { + if next, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID); err != nil || next != cursor { t.Fatal("rejected execution operation published events", next, err) } } diff --git a/services/core/internal/store/executor_principals_test.go b/services/core/internal/store/executor_principals_test.go index 3754e390c..e6921f229 100644 --- a/services/core/internal/store/executor_principals_test.go +++ b/services/core/internal/store/executor_principals_test.go @@ -19,7 +19,7 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.T) { if err != nil || issued.KeyID != keyID || issued.EnvironmentID != "" || len(issued.Token) != 43 { t.Fatal("pre-Session principal issuance failed", err) } - page, err := s.ListSessions(ctx, p.TenantID, "", 10, false, nil) + page, err := sessionAdapter(s).ListSessions(ctx, p.TenantID, "", 10, false, nil) if err != nil || len(page.Sessions) != 0 { t.Fatal("issuance created a Session", err) } @@ -92,7 +92,7 @@ func TestExecutorPrincipalBeforeSessionAndSharedLifecycle(t *testing.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 { + if err := sessionService(t, restarted).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: p.TenantID, SessionID: first.ID}); err != nil { t.Fatal(err) } check(restarted, one.ID, issued, false) diff --git a/services/core/internal/store/export_test.go b/services/core/internal/store/export_test.go index 9be7be169..6e3040cc3 100644 --- a/services/core/internal/store/export_test.go +++ b/services/core/internal/store/export_test.go @@ -25,6 +25,9 @@ func NewTestStore(t *testing.T) (*Store, *pgxpool.Pool) { return testStore(t) } // SessionAdapter is the Session adapter on s's pool and credential key. func SessionAdapter(s *Store) *sessionpg.Store { return sessionAdapter(s) } +// SessionService is the Session service on SessionAdapter(s). +func SessionService(t testing.TB, s *Store) *sessions.Service { return sessionService(t, s) } + // TransitionTurn moves the Turn as the execution owner does, in a Session // transaction on s's writer. func TransitionTurn(ctx context.Context, s *Store, tenant, session, turn string, transition sessions.TurnTransition) (sessions.Turn, error) { diff --git a/services/core/internal/store/function_execution_test.go b/services/core/internal/store/function_execution_test.go index c59f5f3a7..0aca0f12c 100644 --- a/services/core/internal/store/function_execution_test.go +++ b/services/core/internal/store/function_execution_test.go @@ -47,7 +47,7 @@ func functionState(t *testing.T, h *dispatchHarness, count int) sessions.Session t.Helper() deadline := time.Now().Add(20 * time.Second) for { - state, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) + state, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/function_images_native_test.go b/services/core/internal/store/function_images_native_test.go index bedafe512..3a00438fa 100644 --- a/services/core/internal/store/function_images_native_test.go +++ b/services/core/internal/store/function_images_native_test.go @@ -48,7 +48,7 @@ func TestNativeFunctionImagePublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(t, worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/function_inputs_test.go b/services/core/internal/store/function_inputs_test.go index 05384e8ab..5feae8e64 100644 --- a/services/core/internal/store/function_inputs_test.go +++ b/services/core/internal/store/function_inputs_test.go @@ -229,7 +229,7 @@ func TestFunctionInputsRejectInvalidTargetsAndStorageObjects(t *testing.T) { if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "late", []sessions.Input{input}); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal(err) } - current, err := s.GetSession(t.Context(), tenant, session.ID) + current, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || current.LastTurn.ID != turn || current.LastTurn.Status != sessions.TurnFailed { t.Fatal(current, err) } diff --git a/services/core/internal/store/function_item_events_test.go b/services/core/internal/store/function_item_events_test.go index d6a0b3e4e..0b049b1a1 100644 --- a/services/core/internal/store/function_item_events_test.go +++ b/services/core/internal/store/function_item_events_test.go @@ -21,7 +21,7 @@ func TestFunctionResultEventsAreInputs(t *testing.T) { if err := sessionExecution(t, executionWriter(t, s).lease).AppendTurnEvents(t.Context(), tenant, session.ID, turn, 1, events); err != nil { t.Fatal(err) } - changes, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + changes, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil { t.Fatal(err) } @@ -109,7 +109,7 @@ func TestFunctionResultItemsRetainSubmittedFields(t *testing.T) { results++ } } - changes, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + changes, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/function_public_native_test.go b/services/core/internal/store/function_public_native_test.go index 1df47e840..87c754e53 100644 --- a/services/core/internal/store/function_public_native_test.go +++ b/services/core/internal/store/function_public_native_test.go @@ -76,7 +76,7 @@ func nativePublicFunctionServer(t *testing.T, h *dispatchHarness, ctx context.Co }) token := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}}) - handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(worker), nativeDeploymentDefaults("gpt-5.5", provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(t, worker), nativeDeploymentDefaults("gpt-5.5", provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/function_state_test.go b/services/core/internal/store/function_state_test.go index 24da5e4c5..320650e7f 100644 --- a/services/core/internal/store/function_state_test.go +++ b/services/core/internal/store/function_state_test.go @@ -17,7 +17,7 @@ func TestFunctionStateSnapshotsRecoveryAndRetries(t *testing.T) { tenant, session := newTurnSession(t, s) turn := submitMessage(t, s, tenant, session.ID, "start").TurnID transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) - before, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + before, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -53,7 +53,7 @@ func TestFunctionStateSnapshotsRecoveryAndRetries(t *testing.T) { } assertFunctionState(t, s, tenant, session.ID, status, 1-i) } - changes, err := s.ListSessionEvents(t.Context(), tenant, session.ID, before) + changes, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, before) if err != nil { t.Fatal(err) } @@ -76,10 +76,10 @@ func TestFunctionStateSnapshotsRecoveryAndRetries(t *testing.T) { } } } - if _, err := s.GetSession(t.Context(), uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSession(t.Context(), uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - if _, err := s.ListSessionEvents(t.Context(), uuid.NewString(), session.ID, before); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).ListSessionEvents(t.Context(), uuid.NewString(), session.ID, before); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } } @@ -96,14 +96,14 @@ func TestFunctionStateCancellationAndTerminalCleanup(t *testing.T) { t.Fatal(err) } if status == sessions.TurnCancelled { - before, _ := s.SessionEventCursor(t.Context(), tenant, session.ID) + before, _ := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) for _, key := range []string{"cancel", "cancel", "another-cancel"} { if _, err := s.RequestCancel(t.Context(), tenant, session.ID, key); err != nil { t.Fatal(err) } } assertFunctionState(t, s, tenant, session.ID, sessions.TurnWaiting, 0) - changes, err := s.ListSessionEvents(t.Context(), tenant, session.ID, before) + changes, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, before) if err != nil || len(changes) != 1 || changes[0].Event.Type != "agent.session.in_progress" || len(changes[0].RequiredActions) != 0 || changes[0].Turn.CancelRequestedAt.IsZero() { t.Fatal(changes, err) } @@ -154,7 +154,7 @@ func TestFunctionStateReadsRemainConsistentDuringReceipts(t *testing.T) { }) defer wg.Wait() for { - current, err := s.GetSession(t.Context(), tenant, session.ID) + current, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -171,11 +171,11 @@ func TestFunctionStateReadsRemainConsistentDuringReceipts(t *testing.T) { func assertFunctionState(t *testing.T, s *Store, tenant, sessionID, status string, count int) { t.Helper() - current, err := s.GetSession(t.Context(), tenant, sessionID) + current, err := sessionAdapter(s).GetSession(t.Context(), tenant, sessionID) if err != nil || current.LastTurn == nil || current.LastTurn.Status != status || len(current.RequiredActions) != count { t.Fatalf("state: %+v; %v", current, err) } - page, err := s.ListSessions(t.Context(), tenant, "", 10, true, nil) + page, err := sessionAdapter(s).ListSessions(t.Context(), tenant, "", 10, true, nil) if err != nil || len(page.Sessions) != 1 || len(page.Sessions[0].RequiredActions) != count || page.Sessions[0].LastTurn.Status != status { t.Fatal("list differs from retrieve", page, err) } diff --git a/services/core/internal/store/harness_onboarding_test.go b/services/core/internal/store/harness_onboarding_test.go index 60879b0ff..3203b07de 100644 --- a/services/core/internal/store/harness_onboarding_test.go +++ b/services/core/internal/store/harness_onboarding_test.go @@ -63,7 +63,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { }() token := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}}) - handler, err := publicHandler(t, h.s, h.db, auth, "fixture_harness", workerExecution(worker), withPolicy(policy)) + handler, err := publicHandler(t, h.s, h.db, auth, "fixture_harness", workerExecution(t, worker), withPolicy(policy)) if err != nil { t.Fatal(err) } @@ -90,7 +90,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { if err = json.Unmarshal(res.Body.Bytes(), &created); err != nil || created.ID == "" { t.Fatal(res.Body, err) } - h.session, err = h.s.GetSession(ctx, h.tenant, created.ID) + h.session, err = store.SessionAdapter(h.s).GetSession(ctx, h.tenant, created.ID) if err != nil { t.Fatal(err) } @@ -152,7 +152,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { t.Fatalf("incapable runtime received work: %s", p.RunID) case <-time.After(700 * time.Millisecond): } - queued, err := h.s.GetSession(ctx, h.tenant, created.ID) + queued, err := store.SessionAdapter(h.s).GetSession(ctx, h.tenant, created.ID) if err != nil || queued.LastTurn == nil || queued.LastTurn.Status != sessions.TurnQueued { t.Fatal(queued, err) } diff --git a/services/core/internal/store/hosted_initialization_failure_public_test.go b/services/core/internal/store/hosted_initialization_failure_public_test.go index ebb1dd329..90ccc4f5a 100644 --- a/services/core/internal/store/hosted_initialization_failure_public_test.go +++ b/services/core/internal/store/hosted_initialization_failure_public_test.go @@ -193,15 +193,15 @@ func TestHostedInitializationFailureRecordsSafeSessionFailure(t *testing.T) { t.Fatal("failed initialization continued or reclaimed compute", p.steps, p.kills) } - read, err := s.GetSession(t.Context(), tenant, session.ID) + read, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || read.Environment.Status != "failed" || read.EnvironmentFailure == nil || read.EnvironmentFailure.Reason != test.reason || read.EnvironmentInputActivity != nil || read.LastTurn != nil { t.Fatal("Session read", read.EnvironmentFailure, err) } - page, err := s.ListSessions(t.Context(), tenant, "", 10, false, nil) + page, err := store.SessionAdapter(s).ListSessions(t.Context(), tenant, "", 10, false, nil) if err != nil || len(page.Sessions) != 1 || !reflect.DeepEqual(page.Sessions[0].EnvironmentFailure, read.EnvironmentFailure) { t.Fatal("Session list", page, err) } - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + events, err := store.SessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) // Transport may connect while initialization is still running. if len(events) > 0 && events[0].Event.Type == "agent.session.environment.connected" { events = events[1:] @@ -232,13 +232,13 @@ func TestHostedInitializationFailureRecordsSafeSessionFailure(t *testing.T) { } // Tenant B cannot observe the failure. other := uuid.NewString() - if _, err := s.GetSession(t.Context(), other, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := store.SessionAdapter(s).GetSession(t.Context(), other, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign Session read", err) } - if _, err := s.ListSessionEvents(t.Context(), other, session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { + if _, err := store.SessionAdapter(s).ListSessionEvents(t.Context(), other, session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign Session events", err) } - if page, err := s.ListSessions(t.Context(), other, "", 10, false, nil); err != nil || len(page.Sessions) != 0 { + if page, err := store.SessionAdapter(s).ListSessions(t.Context(), other, "", 10, false, nil); err != nil || len(page.Sessions) != 0 { t.Fatal("foreign Session list", page, err) } }) @@ -257,12 +257,12 @@ func TestHostedInitializationFailureSettlesPendingInitialInput(t *testing.T) { p := &hostedFailureProvider{lifecycleProvider: lifecycleProvider{resources: map[string]sandbox.Info{}}, fail: "setup", result: failedInitialization(3)} failHostedInitialization(t, s, db, tenant, environment, p) - read, err := s.GetSession(t.Context(), tenant, session.ID) + read, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || read.PendingInput || read.EnvironmentInputActivity == nil || read.EnvironmentInputActivity.Status != "failed" || read.EnvironmentInputActivity.Failure != "environment_unavailable" || read.EnvironmentFailure == nil { t.Fatal("pending input settlement", read.EnvironmentInputActivity, read.EnvironmentFailure, err) } - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + events, err := store.SessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil { t.Fatal(err) } @@ -300,7 +300,7 @@ func TestHostedInitializationFailurePublicHTTP(t *testing.T) { {OrganizationID: "test-org", ProjectID: tenant, SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "tenant-b", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(w)) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, w)) if err != nil { t.Fatal(err) } @@ -341,7 +341,7 @@ func TestHostedInitializationFailurePublicHTTP(t *testing.T) { awaitInitialization(t, db.pool, tenant, environment.ID, "failed") reason := `Failed to provision environment: script "setup_commands[0]" failed with exit code 3` - read, err := s.GetSession(t.Context(), tenant, session.ID) + read, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || read.EnvironmentFailure == nil { t.Fatal(err) } diff --git a/services/core/internal/store/item_serialization_test.go b/services/core/internal/store/item_serialization_test.go index a74bd030c..828b6b30f 100644 --- a/services/core/internal/store/item_serialization_test.go +++ b/services/core/internal/store/item_serialization_test.go @@ -18,7 +18,7 @@ type wireEvent struct { // wireEvents renders recorded Session events as clients receive them. func wireEvents(t *testing.T, s *Store, tenant, session string, after int64) []wireEvent { t.Helper() - changes, err := s.ListSessionEvents(t.Context(), tenant, session, after) + changes, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session, after) if err != nil { t.Fatal(err) } @@ -72,7 +72,7 @@ func TestAssistantMessageEventsFollowOfficialSequence(t *testing.T) { tenant, session := newTurnSession(t, s) turn := submitMessage(t, s, tenant, session.ID, "start").TurnID transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) - cursor, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -127,7 +127,7 @@ func TestAssistantMessageEventsFollowOfficialSequence(t *testing.T) { func TestInputItemEventsCarryNullOutputIndexAndPhase(t *testing.T) { s, _ := testStore(t) tenant, session := newTurnSession(t, s) - cursor, err := s.SessionEventCursor(t.Context(), tenant, session.ID) + cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/list_query_public_test.go b/services/core/internal/store/list_query_public_test.go index 72bcf988b..3f53a6b23 100644 --- a/services/core/internal/store/list_query_public_test.go +++ b/services/core/internal/store/list_query_public_test.go @@ -42,7 +42,7 @@ func TestListQueryOfficialClientPostgres(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(worker)) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, worker)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/local_environment_devices_test.go b/services/core/internal/store/local_environment_devices_test.go index a16b46b3d..64ad74210 100644 --- a/services/core/internal/store/local_environment_devices_test.go +++ b/services/core/internal/store/local_environment_devices_test.go @@ -59,7 +59,7 @@ func TestEnvironmentDeviceAuthorityAndLifecycle(t *testing.T) { if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), bound.ID); err != nil || !ok { t.Fatalf("valid credential unavailable: %v", err) } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), bound.ID); err != nil || ok { 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 1f255c6e7..ed77c1edf 100644 --- a/services/core/internal/store/local_environment_file_write_test.go +++ b/services/core/internal/store/local_environment_file_write_test.go @@ -79,7 +79,7 @@ func TestLocalEnvironmentFileWriteOwnsMutationBeforeDispatch(t *testing.T) { 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) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil || session.LastTurn != nil { t.Fatal("upload created model execution", err) } diff --git a/services/core/internal/store/local_environment_worker_test.go b/services/core/internal/store/local_environment_worker_test.go index 001bec8cb..cd8e2429a 100644 --- a/services/core/internal/store/local_environment_worker_test.go +++ b/services/core/internal/store/local_environment_worker_test.go @@ -81,7 +81,7 @@ func TestLocalEnvironmentWorkerDirectoryUsesExactAuthorityWithoutModel(t *testin if got := awaitDirectoryResult(t, result); got.err != nil || len(got.value.Entries) != 1 { t.Fatal("local directory failed", got.err) } - session, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil || session.LastTurn != nil || session.EnvironmentInputActivity != nil { t.Fatal("local read admitted execution", err) } @@ -126,7 +126,7 @@ func TestLocalEnvironmentWorkerSchedulesPreparationWithoutRemoteResolver(t *test if frame.DecodePayload(&prepare) != nil || prepare.Configuration.LocalEnvironment == nil || prepare.Configuration.LocalEnvironment.ID != environment.ID { t.Fatal("local preparation lost identity") } - before, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) + before, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil || before.LastTurn != nil { t.Fatal("preparation admitted execution before readiness", err) } diff --git a/services/core/internal/store/mcode_public_native_test.go b/services/core/internal/store/mcode_public_native_test.go index 2ffb143e3..9cd0dcd28 100644 --- a/services/core/internal/store/mcode_public_native_test.go +++ b/services/core/internal/store/mcode_public_native_test.go @@ -48,7 +48,7 @@ func TestNativeMCodePublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, "mcode", workerExecution(worker), acceptUnavailable(t), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "mcode", workerExecution(t, worker), acceptUnavailable(t), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/message_image_admission_test.go b/services/core/internal/store/message_image_admission_test.go index 9ee2a74a6..24f2b03e9 100644 --- a/services/core/internal/store/message_image_admission_test.go +++ b/services/core/internal/store/message_image_admission_test.go @@ -46,7 +46,7 @@ func TestUnqualifiedImageAdmissionIsAtomic(t *testing.T) { if _, err := worker.SubmitInputs(t.Context(), h.tenant, session.ID, "image-batch", imageAdmissionBatch()); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal("batch accepted an unqualified image", err) } - session, err = h.s.GetSession(t.Context(), h.tenant, session.ID) + session, err = store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, session.ID) if err != nil || session.LastTurn != nil || session.EnvironmentInputActivity != nil { t.Fatal("rejected batch persisted execution activity", err) } diff --git a/services/core/internal/store/message_images_native_test.go b/services/core/internal/store/message_images_native_test.go index 46c2583cb..40019ebc2 100644 --- a/services/core/internal/store/message_images_native_test.go +++ b/services/core/internal/store/message_images_native_test.go @@ -49,7 +49,7 @@ func TestNativeMessageImagePublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(t, worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/model_protocol_native_test.go b/services/core/internal/store/model_protocol_native_test.go index dceecd9bf..d0ee2dbd2 100644 --- a/services/core/internal/store/model_protocol_native_test.go +++ b/services/core/internal/store/model_protocol_native_test.go @@ -79,7 +79,7 @@ func TestNativeModelProtocolPublicExecution(t *testing.T) { token := uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}}) providerRevision := uuid.New() - handler, err := publicHandler(t, h.s, h.db, auth, options.Engine, workerExecution(worker), withPolicy(h.d.Policy), modelProviderDefaults(func(context.Context, string) (*modelconfiguration.Snapshot, error) { + handler, err := publicHandler(t, h.s, h.db, auth, options.Engine, workerExecution(t, worker), withPolicy(h.d.Policy), modelProviderDefaults(func(context.Context, string) (*modelconfiguration.Snapshot, error) { return &modelconfiguration.Snapshot{Model: options.Model, HarnessConfig: options.HarnessConfig, Provider: &options.Provider, Revision: providerRevision}, nil })) if err != nil { diff --git a/services/core/internal/store/native_public_execution_test.go b/services/core/internal/store/native_public_execution_test.go index 496914f2e..8128b85b6 100644 --- a/services/core/internal/store/native_public_execution_test.go +++ b/services/core/internal/store/native_public_execution_test.go @@ -38,7 +38,7 @@ func verifyNativePublicExecution(t *testing.T, h *dispatchHarness, parent contex }() token, foreign := uuid.NewString(), uuid.NewString() auth := newTestAuthenticator(t, []testAPIKey{{OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test-org", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "test-runner", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}}) - handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(worker), nativeDeploymentDefaults("gpt-5.5", provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "codex", workerExecution(t, worker), nativeDeploymentDefaults("gpt-5.5", provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/prepared_dispatch_failure_test.go b/services/core/internal/store/prepared_dispatch_failure_test.go index 23f05ff48..aef19d944 100644 --- a/services/core/internal/store/prepared_dispatch_failure_test.go +++ b/services/core/internal/store/prepared_dispatch_failure_test.go @@ -32,7 +32,7 @@ func TestPreparedDispatchSettlesOnlyReadyInput(t *testing.T) { t.Fatal(err) } case "delete": - if err := h.s.DeleteSession(t.Context(), h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := store.SessionService(t, h.s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("pending input deleted", err) } if err := h.s.CommitLegacyDeletion(t.Context(), h.tenant, h.session.ID); err != nil { diff --git a/services/core/internal/store/prepared_dispatch_test.go b/services/core/internal/store/prepared_dispatch_test.go index 109b662bf..4c30eadc6 100644 --- a/services/core/internal/store/prepared_dispatch_test.go +++ b/services/core/internal/store/prepared_dispatch_test.go @@ -81,7 +81,7 @@ func TestPreparedDispatchPromotesOriginalBatchAndPersistsCompletion(t *testing.T if frame.DecodePayload(&prepare) != nil || len(prepare.Configuration.Input) != 0 || prepare.Configuration.RunID != "" || prepare.Configuration.ConversationID != "" || prepare.Configuration.LocalEnvironment == nil || prepare.Configuration.LocalEnvironment.ID != h.device.EnvironmentID || prepare.Configuration.DisableExecutionEnvironment { t.Fatal("invalid preparation configuration", prepare) } - session, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil || session.LastTurn != nil { t.Fatal("preparation created work before readiness", session, err) } diff --git a/services/core/internal/store/public_execution_test.go b/services/core/internal/store/public_execution_test.go index 8d3301506..6b6d99340 100644 --- a/services/core/internal/store/public_execution_test.go +++ b/services/core/internal/store/public_execution_test.go @@ -68,7 +68,7 @@ func TestExecutionWorkerAdmissionBindingAndRecovery(t *testing.T) { if err != nil || bound.ID != h.device.ID { t.Fatal(bound, err) } - active, err := h.s.GetSession(ctx, h.tenant, h.session.ID) + active, err := store.SessionAdapter(h.s).GetSession(ctx, h.tenant, h.session.ID) if err != nil || active.LastTurn == nil || active.LastTurn.Status != sessions.TurnInProgress { t.Fatal(active, err) } @@ -138,7 +138,7 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { if !ended { want = nil } - session, err := h.s.GetSession(ctx, h.tenant, h.session.ID) + session, err := store.SessionAdapter(h.s).GetSession(ctx, h.tenant, h.session.ID) if err != nil || string(session.Usage) != string(want) { t.Fatalf("Session and Turn measurement differ: %+v %v", session, err) } @@ -154,18 +154,18 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { if err := worker.Run(stopped); err != context.Canceled { t.Fatal(err) } - interrupted, err := h.s.GetSession(ctx, h.tenant, h.session.ID) + interrupted, err := store.SessionAdapter(h.s).GetSession(ctx, h.tenant, h.session.ID) if err != nil || interrupted.LastTurn.Status != sessions.TurnFailed { t.Fatal(interrupted, err) } - pending, err := h.s.GetSession(ctx, h.tenant, queued.ID) + pending, err := store.SessionAdapter(h.s).GetSession(ctx, h.tenant, queued.ID) if err != nil || pending.LastTurn.Status != sessions.TurnQueued { t.Fatal(pending, err) } if _, err := h.s.RequestCancel(ctx, h.tenant, queued.ID, "stop-before-dispatch"); err != nil { t.Fatal(err) } - pending, err = h.s.GetSession(ctx, h.tenant, queued.ID) + pending, err = store.SessionAdapter(h.s).GetSession(ctx, h.tenant, queued.ID) if err != nil || pending.LastTurn.Status != sessions.TurnCancelled { t.Fatal(pending, err) } diff --git a/services/core/internal/store/public_handler_fixture_test.go b/services/core/internal/store/public_handler_fixture_test.go index 111eeb790..7db484e46 100644 --- a/services/core/internal/store/public_handler_fixture_test.go +++ b/services/core/internal/store/public_handler_fixture_test.go @@ -85,15 +85,16 @@ func publicHandler(t testing.TB, s *store.Store, db fixtureDB, keys fixtureKeyRe EnvironmentTemplates: environmentTemplates, EnvironmentTemplatesReader: templates, Skills: skillService, SkillsReader: skillStore, Agents: agentService, AgentsReader: agentStore, - Sessions: s, + Sessions: sessionService, + SessionsReader: sessionStore, SessionCreation: s, - SessionEvents: s, + SessionEvents: sessionStore, Turns: sessionStore, Items: sessionStore, Subagents: sessionStore, Artifacts: sessionService, ArtifactsReader: sessionStore, - SessionAdmin: s, Environments: sessionService, EnvironmentsReader: sessionStore, Admin: s, AdminRuntimeTargets: sessionStore, AdminAudit: audit, WriteAudit: audit, + SessionAdmin: sessionStore, Environments: sessionService, EnvironmentsReader: sessionStore, Admin: s, AdminRuntimeTargets: sessionStore, AdminAudit: audit, WriteAudit: audit, ExecutorConnections: strict, Metrics: strict, RuntimeObservations: strict, RuntimeHistory: strict, } for _, c := range configure { @@ -158,14 +159,15 @@ func storeExecution(t testing.TB, s *store.Store) func(*api.Dependencies) { } } -// workerExecution runs Sessions through worker. -func workerExecution(worker *execution.Worker) func(*api.Dependencies) { +// workerExecution runs Sessions through worker. It wires no archive; an +// archive request fails the test. +func workerExecution(t testing.TB, worker *execution.Worker) func(*api.Dependencies) { return func(d *api.Dependencies) { d.Execution = &api.Execution{ ExecutorURL: testExecutorURL, SessionAdmission: worker, InputAdmission: worker, - SessionArchive: worker, + SessionArchive: strictStandIn{t}, Workspaces: worker, } } @@ -261,8 +263,8 @@ func (s strictStandIn) ExecutorConnected(context.Context, string, string) (bool, return false, nil } -func (s strictStandIn) ArchiveManagedSession(context.Context, string, string, uint64) (sessions.ManagedArchive, error) { - s.unexpected("ArchiveManagedSession") +func (s strictStandIn) ArchiveSession(context.Context, string, string, uint64) (sessions.ManagedArchive, error) { + s.unexpected("ArchiveSession") return sessions.ManagedArchive{}, nil } diff --git a/services/core/internal/store/runtime_allocations_test.go b/services/core/internal/store/runtime_allocations_test.go index 9ee59b425..acf21784f 100644 --- a/services/core/internal/store/runtime_allocations_test.go +++ b/services/core/internal/store/runtime_allocations_test.go @@ -52,7 +52,7 @@ func TestRuntimeAllocationAtomicOwnershipAndRecovery(t *testing.T) { if err != nil || !ok || credential.CredentialHash != runtimedevice.HashCredential(secret) { t.Fatal("retry rewrote bootstrap credential") } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } retained, err := deploymentStore(next).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: environment.ID}) diff --git a/services/core/internal/store/runtime_compute_lifecycle_test.go b/services/core/internal/store/runtime_compute_lifecycle_test.go index 4e175095b..f32f88f52 100644 --- a/services/core/internal/store/runtime_compute_lifecycle_test.go +++ b/services/core/internal/store/runtime_compute_lifecycle_test.go @@ -393,7 +393,7 @@ func TestRuntimeComputeLifecycleIdleSuspendAndQueuedSameSessionWake(t *testing.T if err := f.db.pool.QueryRow(t.Context(), `SELECT count(*) FILTER(WHERE id=$2 AND status='completed'),count(*) FILTER(WHERE id=$3 AND status='queued') FROM turns WHERE session_id=$1`, session.ID, completed, queued).Scan(&completedCount, &queuedCount); err != nil || completedCount != 1 || queuedCount != 1 { t.Fatal("wake replayed/consumed prior or next Turn", err) } - if got, err := f.store.GetSession(t.Context(), tenant, session.ID); err != nil || string(got.Configuration) != string(session.Configuration) { + if got, err := store.SessionAdapter(f.store).GetSession(t.Context(), tenant, session.ID); err != nil || string(got.Configuration) != string(session.Configuration) { t.Fatal("configuration changed during restore", err) } if f.provider.promptFrames.Load() != 0 { @@ -510,7 +510,7 @@ func TestRuntimeComputeLifecycleSuspendedDeletionAndExpiryCleanup(t *testing.T) f.complete(owner) f.phase(tenant, env.ID, "suspended") if kind == "deleted" { - if err := f.store.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := store.SessionService(t, f.store).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } } else { diff --git a/services/core/internal/store/runtime_configuration_cleanup_test.go b/services/core/internal/store/runtime_configuration_cleanup_test.go index 4911133ca..855fc427c 100644 --- a/services/core/internal/store/runtime_configuration_cleanup_test.go +++ b/services/core/internal/store/runtime_configuration_cleanup_test.go @@ -6,6 +6,8 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -97,7 +99,7 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { t.Fatal("unexpected creation outcome", owner, err) } if test.rejectCreate { - value, err := s.GetSession(t.Context(), tenant, session.ID) + value, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || value.Environment.Status != "failed" { t.Fatal("rejected configuration did not terminate its retained Session", value, err) } @@ -119,7 +121,7 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { } // Exercise the existing authorized deletion lifecycle; no new // deployment cleanup API is introduced by this regression test. - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } } @@ -152,7 +154,7 @@ func TestManagedRuntimeConfigurationCleanup(t *testing.T) { } if test.rejectCreate { // The rejected Session remains publicly readable after cleanup. - if _, err := s.GetSession(t.Context(), tenant, session.ID); err != nil { + if _, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, session.ID); err != nil { t.Fatal("configuration rejection deleted Session history", err) } } diff --git a/services/core/internal/store/runtime_connection_test.go b/services/core/internal/store/runtime_connection_test.go index de89b8fe9..e40311ecd 100644 --- a/services/core/internal/store/runtime_connection_test.go +++ b/services/core/internal/store/runtime_connection_test.go @@ -18,6 +18,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtime" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { @@ -87,7 +88,7 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { } defer conn.Close() assertStatus("connected") - got, err := s.GetSession(t.Context(), tenant, session.ID) + got, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || got.LastTurn != nil || got.EnvironmentInputActivity != nil { t.Fatal("connection fabricated native execution", err) } @@ -111,7 +112,7 @@ func TestManagedRuntimeConnectionTracksAuthenticatedSocket(t *testing.T) { if err != nil || retained.ID != owner.ID || retained.DeviceID != owner.DeviceID || p.creates != 1 { t.Fatal("restart replaced Runtime identity", err) } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } reconcileManagedState(t, w, db, tenant, environment.ID, "released") diff --git a/services/core/internal/store/runtime_creation_settlement_test.go b/services/core/internal/store/runtime_creation_settlement_test.go index 008f31f04..5b812681a 100644 --- a/services/core/internal/store/runtime_creation_settlement_test.go +++ b/services/core/internal/store/runtime_creation_settlement_test.go @@ -7,6 +7,8 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -64,7 +66,7 @@ func TestManagedRuntimeConfirmedAbsentCreateReleasesAtomically(t *testing.T) { if _, ok, err := fixtureSessionStore(db).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("released credential retained authority", err) } - value, err := s.GetSession(t.Context(), tenant, session.ID) + value, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || value.Environment.Status != "failed" { t.Fatal("environment not terminated", err) } @@ -100,7 +102,7 @@ func TestManagedRuntimeObservedSettlementAllowsOwnedCleanup(t *testing.T) { if _, err := w.ProvisionEnvironment(t.Context(), tenant, environment.ID, key); err == nil { t.Fatal("uncertain Create succeeded") } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } reconcileManagedState(t, w, db, tenant, environment.ID, "released") diff --git a/services/core/internal/store/runtime_deployment_test.go b/services/core/internal/store/runtime_deployment_test.go index 6868bb647..8a2f82cbf 100644 --- a/services/core/internal/store/runtime_deployment_test.go +++ b/services/core/internal/store/runtime_deployment_test.go @@ -62,7 +62,7 @@ func TestRuntimeDeploymentPendingSessionsCannotMigrate(t *testing.T) { if err := deploymentExecution(t, w).ConfigureProcess(t.Context(), nil); err == nil { t.Fatal("pending Session orphaned by removing provider") } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } deploymentConfigure(t, w, &next) @@ -81,7 +81,7 @@ func TestRuntimeDeploymentUnknownAllocationsBlockAdoptionAndSwitch(t *testing.T) if err := deploymentExecution(t, w).ConfigureProcess(t.Context(), &old); err == nil || !strings.Contains(err.Error(), "no verified backend identity") { t.Fatal("legacy allocation silently adopted", err) } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } if _, err := deploymentExecution(t, w).RequestCleanup(t.Context(), owner); err != nil { diff --git a/services/core/internal/store/runtime_enrollment_connection_test.go b/services/core/internal/store/runtime_enrollment_connection_test.go index fde167d43..f7cebbbda 100644 --- a/services/core/internal/store/runtime_enrollment_connection_test.go +++ b/services/core/internal/store/runtime_enrollment_connection_test.go @@ -168,7 +168,7 @@ func TestEnrolledDaemonConnectionRevocationAndRestart(t *testing.T) { if err = db.pool.QueryRow(t.Context(), "SELECT count(*) FROM runtime_allocations WHERE environment_id=$1", environment.ID).Scan(&allocations); err != nil || allocations != 0 { t.Fatal("user Runtime acquired managed allocation", allocations, err) } - current, err := s.GetSession(t.Context(), principal.TenantID, session.ID) + current, err := store.SessionAdapter(s).GetSession(t.Context(), principal.TenantID, session.ID) if err != nil || current.LastTurn != nil { t.Fatal("connection handling created execution", err) } diff --git a/services/core/internal/store/runtime_enrollment_test.go b/services/core/internal/store/runtime_enrollment_test.go index ab2e10932..bf2e20d87 100644 --- a/services/core/internal/store/runtime_enrollment_test.go +++ b/services/core/internal/store/runtime_enrollment_test.go @@ -137,7 +137,7 @@ func TestRuntimeEnrollmentConcurrentAndDeletion(t *testing.T) { if err := pool.QueryRow(t.Context(), "SELECT count(*) FROM runtime_allocations WHERE environment_id=$1", environment.ID).Scan(&allocations); err != nil || allocations != 0 { t.Fatalf("enrollment allocated compute: %d %v", allocations, err) } - if err := s.DeleteSession(t.Context(), p.TenantID, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: p.TenantID, SessionID: session.ID}); err != nil { t.Fatal(err) } if _, err := sessionService(t, s).EnrollRuntime(t.Context(), environment.ID, executorDigest(key.Token)); !errors.Is(err, sessions.ErrNotFound) { diff --git a/services/core/internal/store/runtime_environment_terminal_test.go b/services/core/internal/store/runtime_environment_terminal_test.go index c22bfed18..3a1ed2d7f 100644 --- a/services/core/internal/store/runtime_environment_terminal_test.go +++ b/services/core/internal/store/runtime_environment_terminal_test.go @@ -38,7 +38,7 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin if _, err := deploymentExecution(t, writer).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } - ended, err := s.GetSession(t.Context(), tenant, session.ID) + ended, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) status := "failed" if expired { status = "expired" @@ -63,7 +63,7 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin if _, err := s.CreateSession(t.Context(), tenant, input); err != nil { t.Fatal("matching creation retry changed outcome", err) } - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + events, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil { t.Fatal(err) } @@ -93,11 +93,11 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin } else if last.EnvironmentFailure != nil { t.Fatal("expiry recorded a provisioning failure", last) } - cursor, _ := s.SessionEventCursor(t.Context(), tenant, session.ID) + cursor, _ := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) if _, err := deploymentExecution(t, writer).RequestCleanup(t.Context(), owner); err != nil { t.Fatal(err) } - if next, err := s.SessionEventCursor(t.Context(), tenant, session.ID); err != nil || next != cursor { + if next, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID); err != nil || next != cursor { t.Fatal("cleanup repeated terminal events", next, err) } if err := sessionExecution(t, writer.lease).ReplaceEnvironmentConnection(t.Context(), tenant, session.Environment.ID, uuid.NewString()); !errors.Is(err, sessions.ErrInvalidInput) { @@ -135,7 +135,7 @@ func TestManagedEnvironmentFailureRollsBackWithSessionEvent(t *testing.T) { if _, err := deploymentExecution(t, writer).RequestCleanup(t.Context(), owner); err == nil { t.Fatal("cleanup committed without failure event") } - current, err := s.GetSession(t.Context(), tenant, session.ID) + current, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || current.Environment.Status != "pending" || current.EnvironmentInputActivity != nil { t.Fatal("partial public termination", err) } diff --git a/services/core/internal/store/runtime_lifecycle_nodes_test.go b/services/core/internal/store/runtime_lifecycle_nodes_test.go index 248b7645f..f7c8fea32 100644 --- a/services/core/internal/store/runtime_lifecycle_nodes_test.go +++ b/services/core/internal/store/runtime_lifecycle_nodes_test.go @@ -158,7 +158,7 @@ func TestRuntimeLifecycleNodeInventoryAndRouting(t *testing.T) { if _, err := s.pool.Exec(t.Context(), "UPDATE runtime_allocations SET node_id=$1 WHERE id=$2", other, owner.ID); err != nil { t.Fatal(err) } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } checkRoute(other, nil) // Deletion does not discard cleanup routing. diff --git a/services/core/internal/store/runtime_lifecycle_test.go b/services/core/internal/store/runtime_lifecycle_test.go index 1473b627a..b242d2d32 100644 --- a/services/core/internal/store/runtime_lifecycle_test.go +++ b/services/core/internal/store/runtime_lifecycle_test.go @@ -170,7 +170,7 @@ func TestManagedRuntimeLostCreateRestartAndDeletion(t *testing.T) { if err != nil || !retry.Replayed || retry.ID != owner.ID || p.creates != 1 { t.Fatal("restart replayed Create") } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } // A scan may first exhaust its previous cursor before starting a new cycle. @@ -194,7 +194,7 @@ func TestManagedRuntimeUnknownCreationRetainsCleanup(t *testing.T) { if err == nil { t.Fatal("expected uncertain creation") } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } reconcileManagedState(t, w, db, tenant, env.ID, "cleanup_pending") diff --git a/services/core/internal/store/runtime_node_lifecycle_test.go b/services/core/internal/store/runtime_node_lifecycle_test.go index 7055fe175..d0b6e26e9 100644 --- a/services/core/internal/store/runtime_node_lifecycle_test.go +++ b/services/core/internal/store/runtime_node_lifecycle_test.go @@ -9,6 +9,8 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -100,7 +102,7 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { if err := f.nodes.TouchActivity(t.Context(), wakeTenant, wakeEnv.ID); err != nil { t.Fatal(err) } - if err := f.store.DeleteSession(t.Context(), deleteTenant, deleteSession.ID); err != nil { + if err := store.SessionService(t, f.store).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: deleteTenant, SessionID: deleteSession.ID}); err != nil { t.Fatal(err) } @@ -135,7 +137,7 @@ func TestManagedNodesIsolateBlockedProviderAndInitialization(t *testing.T) { f.online(nodeC) ct, cs, ce := f.session(nodeC, false) f.provision(ct, ce) - if err := f.store.DeleteSession(t.Context(), ct, cs.ID); err != nil { + if err := store.SessionService(t, f.store).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: ct, SessionID: cs.ID}); err != nil { t.Fatal(err) } waitNodeIsolation(t, 7*time.Second, func() (bool, string) { diff --git a/services/core/internal/store/runtime_nodes.go b/services/core/internal/store/runtime_nodes.go deleted file mode 100644 index ef13d1565..000000000 --- a/services/core/internal/store/runtime_nodes.go +++ /dev/null @@ -1,13 +0,0 @@ -package store - -import ( - "github.com/google/uuid" - "github.com/jackc/pgx/v5/pgtype" -) - -func runtimeUUID(id pgtype.UUID) string { - if !id.Valid { - return "" - } - return uuid.UUID(id.Bytes).String() -} diff --git a/services/core/internal/store/runtime_nodes_test.go b/services/core/internal/store/runtime_nodes_test.go index 7627a1421..93b247828 100644 --- a/services/core/internal/store/runtime_nodes_test.go +++ b/services/core/internal/store/runtime_nodes_test.go @@ -97,7 +97,7 @@ type sessionPlacement struct { // sessionRuntimePlacement reads the node a Session was placed on. func sessionRuntimePlacement(ctx context.Context, s *Store, tenant, session string) (sessionPlacement, error) { - value, err := s.GetSession(ctx, tenant, session) + value, err := sessionAdapter(s).GetSession(ctx, tenant, session) if err != nil { return sessionPlacement{}, err } @@ -112,7 +112,11 @@ func sessionRuntimePlacement(ctx context.Context, s *Store, tenant, session stri if errors.Is(err, pgx.ErrNoRows) { return sessionPlacement{}, sessions.ErrNotFound } - return sessionPlacement{NodeID: runtimeUUID(p.NodeID), Available: p.Available && !p.ReleasedAt.Valid}, err + placement := sessionPlacement{Available: p.Available && !p.ReleasedAt.Valid} + if p.NodeID.Valid { + placement.NodeID = uuid.UUID(p.NodeID.Bytes).String() + } + return placement, err } func TestRuntimeNodesAtomicPlacementAndRetry(t *testing.T) { s, _, d := managerFixture(t, 1, 4) @@ -157,7 +161,7 @@ func TestRuntimeNodesAtomicPlacementAndRetry(t *testing.T) { if err := service.RemoveNode(t.Context(), d.LocalNodeID); !errors.Is(err, deployment.ErrNodeInUse) { t.Fatal("removed pending placement", err) } - if err := s.DeleteSession(t.Context(), tenant, retained.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: retained.ID}); err != nil { t.Fatal(err) } input := managerSessionInput("retry") @@ -253,10 +257,10 @@ func TestRuntimeNodesRetention(t *testing.T) { if err := nodes.RemoveNode(t.Context(), next.LocalNodeID); !errors.Is(err, deployment.ErrNodeInUse) { t.Fatal(err) } - if err := s.DeleteSession(t.Context(), tenant, pending.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: pending.ID}); err != nil { t.Fatal(err) } - if err := s.DeleteSession(t.Context(), tenant, first.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: first.ID}); err != nil { t.Fatal(err) } retained, err = deploymentExecution(t, w).RequestCleanup(t.Context(), retained) diff --git a/services/core/internal/store/runtime_observation_scan_test.go b/services/core/internal/store/runtime_observation_scan_test.go index 9bc3238b3..1e43c760b 100644 --- a/services/core/internal/store/runtime_observation_scan_test.go +++ b/services/core/internal/store/runtime_observation_scan_test.go @@ -3,6 +3,9 @@ package store_test import ( "slices" "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *testing.T) { @@ -13,7 +16,7 @@ func TestRuntimeObservationScanIsDeploymentWideBoundedAndExcludesDeleted(t *test expected = append(expected, session.ID) } slices.Sort(expected) - if err := s.DeleteSession(t.Context(), sessionTenant(t, db, expected[2]), expected[2]); err != nil { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: sessionTenant(t, db, expected[2]), SessionID: expected[2]}); err != nil { t.Fatal(err) } expected = append(expected[:2], expected[3:]...) diff --git a/services/core/internal/store/runtime_pending_test.go b/services/core/internal/store/runtime_pending_test.go index 9a94fefa6..c7b95dea5 100644 --- a/services/core/internal/store/runtime_pending_test.go +++ b/services/core/internal/store/runtime_pending_test.go @@ -22,7 +22,7 @@ func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) t.Fatal(err) } _, deleted, deletedEnvironment := managedSession(t, s, db) - if err := s.DeleteSession(t.Context(), deleted.TenantID, deleted.ID); err != nil { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: deleted.TenantID, SessionID: deleted.ID}); err != nil { t.Fatal(err) } key := uuid.NewString() @@ -60,11 +60,11 @@ func TestManagedRuntimeAutomaticBootstrapRecoversCommittedSessions(t *testing.T) if _, err := fixtureReader(db).EnvironmentAllocation(t.Context(), deployment.AllocationKey{TenantID: deleted.TenantID, EnvironmentID: deletedEnvironment.ID}); err == nil { t.Fatal("deleted Session provisioned") } - waiting, err := s.GetSession(t.Context(), tenant, initial.ID) + waiting, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, initial.ID) if err != nil || waiting.LastTurn != nil || waiting.EnvironmentInputActivity != nil { t.Fatal("compute existence claimed input readiness", waiting, err) } - quiet, err := s.GetSession(t.Context(), tenant, idle.ID) + quiet, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, idle.ID) if err != nil || quiet.LastTurn != nil || quiet.EnvironmentInputActivity != nil { t.Fatal("idle creation fabricated work", err) } diff --git a/services/core/internal/store/runtime_scan_test.go b/services/core/internal/store/runtime_scan_test.go index fe2305218..dd1e2cac6 100644 --- a/services/core/internal/store/runtime_scan_test.go +++ b/services/core/internal/store/runtime_scan_test.go @@ -12,6 +12,8 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) type scanProvider struct { @@ -87,7 +89,7 @@ func TestManagedRuntimeScanEmptyAfterCleanupAndCanceledCall(t *testing.T) { if len(p.observed) != 0 { t.Fatal("canceled scan reached provider") } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } // EOF must not delay cleanup to a second call, and cancellation must diff --git a/services/core/internal/store/runtime_suspension_test.go b/services/core/internal/store/runtime_suspension_test.go index d768c9597..ca07c8f0d 100644 --- a/services/core/internal/store/runtime_suspension_test.go +++ b/services/core/internal/store/runtime_suspension_test.go @@ -181,7 +181,7 @@ func TestRuntimeSuspensionWakeDoesNotLoseNewerWork(t *testing.T) { if _, err := deploymentExecution(t, w).KeepAllocation(t.Context(), owner); err != nil { t.Fatal(err) } - if _, err := s.GetSession(t.Context(), owner.TenantID, owner.SessionID); err != nil { + if _, err := sessionAdapter(s).GetSession(t.Context(), owner.TenantID, owner.SessionID); err != nil { t.Fatal(err) } quiet, err := deploymentStore(w).Activity(t.Context(), owner.ID) @@ -246,7 +246,7 @@ func TestRuntimeSuspensionRetentionAndDeletedSession(t *testing.T) { if _, err := deploymentExecution(t, w).SetCompute(t.Context(), retained, "restoring", json.RawMessage(`{}`), &until, 0); !errors.Is(err, deployment.ErrAllocationConflict) { t.Fatal("expired snapshot restored from stale observation", err) } - if err := s.DeleteSession(t.Context(), owner.TenantID, owner.SessionID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: owner.TenantID, SessionID: owner.SessionID}); err != nil { t.Fatal(err) } if err := deploymentService(t, s).TouchActivity(t.Context(), owner.TenantID, owner.EnvironmentID); !errors.Is(err, sessions.ErrNotFound) { diff --git a/services/core/internal/store/runtime_worker_recovery_test.go b/services/core/internal/store/runtime_worker_recovery_test.go index 56853bb93..50c70e6fe 100644 --- a/services/core/internal/store/runtime_worker_recovery_test.go +++ b/services/core/internal/store/runtime_worker_recovery_test.go @@ -49,7 +49,7 @@ func TestPreparedDispatchKeepsPendingReservationAfterComputeConflict(t *testing. t.Fatal("rejected promotion lost its pending owner", got) } assertPreparationReleased(t, h, frame.ID, handle) - session, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil || session.LastTurn != nil { t.Fatal("blocked promotion started a Turn", session, err) } diff --git a/services/core/internal/store/sandbox_deployment_switch_test.go b/services/core/internal/store/sandbox_deployment_switch_test.go index 42848e08b..2c26d8739 100644 --- a/services/core/internal/store/sandbox_deployment_switch_test.go +++ b/services/core/internal/store/sandbox_deployment_switch_test.go @@ -156,7 +156,7 @@ func TestSandboxDirectDeploymentOwnershipAndCleanSwitch(t *testing.T) { if err := deploymentExecution(t, w).CancelReset(SandboxResetTestContext(t.Context()), id, 1); !errors.Is(err, deployment.ErrConflict) { t.Fatal("stale resume accepted", err) } - if _, err := s.GetSession(t.Context(), tenant, session.ID); err != nil { + if _, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID); err != nil { t.Fatal("historical Session lost", err) } } @@ -355,7 +355,7 @@ func TestSandboxSwitchPreservesReleasedAllocationAndItemHistory(t *testing.T) { t.Fatal("switch rewrote released allocation ownership") } for _, id := range []string{session.ID, history.ID} { - if _, err := s.GetSession(t.Context(), tenant, id); err != nil { + if _, err := sessionAdapter(s).GetSession(t.Context(), tenant, id); err != nil { t.Fatal("switch lost undeleted Session", err) } } diff --git a/services/core/internal/store/sandbox_deployment_switch_worker_test.go b/services/core/internal/store/sandbox_deployment_switch_worker_test.go index 64d6118ff..e2a8ad911 100644 --- a/services/core/internal/store/sandbox_deployment_switch_worker_test.go +++ b/services/core/internal/store/sandbox_deployment_switch_worker_test.go @@ -11,6 +11,7 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/e2b" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" @@ -155,7 +156,7 @@ func TestSandboxWorkerSwitchesAndRecoversFailedActivation(t *testing.T) { if _, err := w.UpdateSandboxDeployment(store.SandboxResetTestContext(t.Context()), next); !errors.Is(err, deployment.ErrConflict) { t.Fatal("dirty switch accepted", err) } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := store.SessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } reconcileManagedState(t, w, db, tenant, environment.ID, "released") diff --git a/services/core/internal/store/sandbox_reset_test.go b/services/core/internal/store/sandbox_reset_test.go index 44a54f2bc..7173b5c3a 100644 --- a/services/core/internal/store/sandbox_reset_test.go +++ b/services/core/internal/store/sandbox_reset_test.go @@ -103,20 +103,20 @@ func TestSandboxResetAutoUsesStartedWorkAndLockedRecheck(t *testing.T) { if err != nil || (len(page) == 0) != busy { t.Fatal("auto eligibility", page, err) } - _, err = w.ArchiveSandboxResetSession(t.Context(), tenant, session.ID, 1, reset.Reset.RequestedAt) + _, err = deploymentExecution(t, w).ArchiveResetSession(t.Context(), tenant, session.ID, 1, reset.Reset.RequestedAt) if busy { - if !errors.Is(err, ErrSandboxResetSessionBusy) { + if !errors.Is(err, deployment.ErrSandboxResetSessionBusy) { t.Fatal("auto cut active work", err) } if err := deploymentExecution(t, w).StartReset(SandboxResetTestContext(t.Context()), installation, deployment.ResetRequest{ExpectedGeneration: 1, Clear: "force"}); err != nil { t.Fatal(err) } - _, err = w.ArchiveSandboxResetSession(t.Context(), tenant, session.ID, 1, reset.Reset.RequestedAt) + _, err = deploymentExecution(t, w).ArchiveResetSession(t.Context(), tenant, session.ID, 1, reset.Reset.RequestedAt) } if err != nil { t.Fatal("archive", err) } - archived, err := s.GetSession(t.Context(), tenant, session.ID) + archived, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || archived.Environment.Status != "expired" { t.Fatal("archive not durable", archived, err) } @@ -151,7 +151,7 @@ func TestSandboxResetCancellationABADeadlineAndGeneration(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err = w.ArchiveSandboxResetSession(t.Context(), tenant, session.ID, 1, first.Reset.RequestedAt); !errors.Is(err, deployment.ErrConflict) { + if _, err = deploymentExecution(t, w).ArchiveResetSession(t.Context(), tenant, session.ID, 1, first.Reset.RequestedAt); !errors.Is(err, deployment.ErrConflict) { t.Fatal("cancelled reset archived successor work", err) } runtimeSuspensionSQL(t, s.pool, `UPDATE runtime_deployment SET reset_deadline_at=clock_timestamp()-interval '1 second'`) @@ -168,7 +168,7 @@ func TestSandboxResetCancellationABADeadlineAndGeneration(t *testing.T) { if err := deploymentExecution(t, w).StartReset(ctx, installation, deployment.ResetRequest{ExpectedGeneration: 0, Clear: "force"}); !errors.Is(err, deployment.ErrConflict) { t.Fatal("stale reset precedence", err) } - if _, err := w.ArchiveSandboxResetSession(t.Context(), tenant, session.ID, 1, second.Reset.RequestedAt); err != nil { + if _, err := deploymentExecution(t, w).ArchiveResetSession(t.Context(), tenant, session.ID, 1, second.Reset.RequestedAt); err != nil { t.Fatal(err) } if _, err := deploymentExecution(t, w).CompleteReset(ctx, installation, 1, first.Reset.RequestedAt); !errors.Is(err, deployment.ErrConflict) { @@ -206,7 +206,7 @@ func TestSandboxResetAutoRechecksTurnStartedAfterListing(t *testing.T) { // Session lock that the later conditional archive must reacquire. turn := submitMessage(t, s, tenant, session.ID, "after-list") transition(t, w, tenant, session.ID, turn.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - if _, err := w.ArchiveSandboxResetSession(t.Context(), tenant, session.ID, 1, reset.Reset.RequestedAt); !errors.Is(err, ErrSandboxResetSessionBusy) { + if _, err := deploymentExecution(t, w).ArchiveResetSession(t.Context(), tenant, session.ID, 1, reset.Reset.RequestedAt); !errors.Is(err, deployment.ErrSandboxResetSessionBusy) { t.Fatal("listed idle candidate cut a new Turn", err) } before := adminMutationSnapshot(t, s, "sessions", "environments", "environment_input_reservations", "turns", "runtime_placements") @@ -261,7 +261,7 @@ func TestSandboxResetSnapshotCountsOfflineOwnershipOnce(t *testing.T) { runtimeSuspensionSQL(t, s.pool, `UPDATE runtime_allocations SET compute_phase='suspended',compute_retained_until=clock_timestamp()+interval '1 hour' WHERE id=$1`, allocation.ID) tenant, deleted := managedArchiveSession(t, s, managerSessionInput(uuid.NewString())) archiveAllocation(t, w, tenant, deleted, process.InstallationID) - if err := s.DeleteSession(t.Context(), tenant, deleted.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: deleted.ID}); err != nil { t.Fatal(err) } reset, err := startReset(t, SandboxResetTestContext(t.Context()), w, process.InstallationID, deployment.ResetRequest{ExpectedGeneration: 1, Clear: "auto"}) @@ -327,7 +327,7 @@ func TestSandboxResetPaginationSkipsBusyPrefixAndPreservesSelfHosted(t *testing. if i <= 34 { runtimeSuspensionSQL(t, s.pool, `INSERT INTO turns(id,session_id,status) VALUES($1,$2,'in_progress')`, uuid.NewString(), id) } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } } @@ -340,7 +340,7 @@ func TestSandboxResetPaginationSkipsBusyPrefixAndPreservesSelfHosted(t *testing. if err != nil || len(page) != 1 || page[0].SessionID != "00000000-0000-4000-8000-000000000035" { t.Fatal("busy prefix starved idle work", page, err) } - if _, err := w.ArchiveSandboxResetSession(t.Context(), page[0].TenantID, page[0].SessionID, 1, reset.Reset.RequestedAt); err != nil { + if _, err := deploymentExecution(t, w).ArchiveResetSession(t.Context(), page[0].TenantID, page[0].SessionID, 1, reset.Reset.RequestedAt); err != nil { t.Fatal(err) } if err := deploymentExecution(t, w).StartReset(SandboxResetTestContext(t.Context()), installation, deployment.ResetRequest{ExpectedGeneration: 1, Clear: "force"}); err != nil { @@ -354,10 +354,10 @@ func TestSandboxResetPaginationSkipsBusyPrefixAndPreservesSelfHosted(t *testing. if err != nil || len(second) != 2 { t.Fatal(second, err) } - if _, err := w.ArchiveSandboxResetSession(t.Context(), selfTenant, self.ID, 1, reset.Reset.RequestedAt); !errors.Is(err, sessions.ErrInvalidInput) { + if _, err := deploymentExecution(t, w).ArchiveResetSession(t.Context(), selfTenant, self.ID, 1, reset.Reset.RequestedAt); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal("self-hosted reset archive", err) } - view, err := s.GetSession(t.Context(), selfTenant, self.ID) + view, err := sessionAdapter(s).GetSession(t.Context(), selfTenant, self.ID) if err != nil || view.Environment.Status == "expired" { t.Fatal("reset changed self-hosted Session", view, err) } @@ -388,7 +388,7 @@ func TestSandboxResetOwnerRestartRetainsDeadlineAndProvenance(t *testing.T) { if err := deploymentExecution(t, successor).AdvanceResetDeadline(t.Context()); err != nil { t.Fatal(err) } - if _, err := successor.ArchiveSandboxResetSession(t.Context(), tenant, session.ID, 1, reset.Reset.RequestedAt); err != nil { + if _, err := deploymentExecution(t, successor).ArchiveResetSession(t.Context(), tenant, session.ID, 1, reset.Reset.RequestedAt); err != nil { t.Fatal(err) } committed, err := deploymentExecution(t, successor).CompleteReset(t.Context(), installation, 1, reset.Reset.RequestedAt) diff --git a/services/core/internal/store/sandbox_specification_store_test.go b/services/core/internal/store/sandbox_specification_store_test.go index 06464c77b..6164daf47 100644 --- a/services/core/internal/store/sandbox_specification_store_test.go +++ b/services/core/internal/store/sandbox_specification_store_test.go @@ -223,7 +223,7 @@ func TestSandboxSpecificationChangesPreserveEveryRetainedResource(t *testing.T) if _, err := deploymentExecution(t, w).ReleaseAllocation(t.Context(), owner); err != nil { t.Fatal(err) } - } else if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + } else if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } changed := input diff --git a/services/core/internal/store/scheduling.go b/services/core/internal/store/scheduling.go index c02516880..b094ed530 100644 --- a/services/core/internal/store/scheduling.go +++ b/services/core/internal/store/scheduling.go @@ -2,15 +2,11 @@ package store import ( "context" - "encoding/json" - "errors" - "fmt" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" - "github.com/jackc/pgx/v5" ) func (s *Store) ListEnvironmentInputWork(ctx context.Context, after string, connectedDevices []string) ([]sessions.EnvironmentInputWork, error) { @@ -28,100 +24,3 @@ func (s *Store) ListEnvironmentInputWork(ctx context.Context, after string, conn } return work, nil } - -func (s *Store) sessionActivity(ctx context.Context, session sessions.Session, err error) (sessions.Session, error) { - if err != nil { - return sessions.Session{}, err - } - err = s.pooled.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { - var err error - session, err = readSessionActivity(ctx, s.queries.WithTx(tx), session) - return err - }) - return session, err -} - -// SessionStreamSnapshot reads the projection that GetSession returns and the -// committed Session event cursor from one database snapshot. -func (s *Store) SessionStreamSnapshot(ctx context.Context, tenantID, sessionID string) (sessions.Session, int64, error) { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.Session{}, 0, err - } - id, err := parseID(sessionID) - if err != nil { - return sessions.Session{}, 0, err - } - var session sessions.Session - var cursor int64 - err = s.pooled.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - row, err := q.GetSession(ctx, sqlc.GetSessionParams{TenantID: tenant, ID: id}) - if err != nil { - return err - } - if session, err = sessionpg.SessionFromRow(row); err != nil { - return err - } - cursor = row.EventSequence - session, err = readSessionActivity(ctx, q, session) - return err - }) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.Session{}, 0, sessions.ErrNotFound - } - if err != nil { - return sessions.Session{}, 0, fmt.Errorf("read session stream snapshot: %w", err) - } - return session, cursor, nil -} - -// readSessionActivity adds the Environment, reservation activity and latest Turn -// projection within the caller's snapshot. -func readSessionActivity(ctx context.Context, q *sqlc.Queries, session sessions.Session) (sessions.Session, error) { - id, _ := parseID(session.ID) - tenant, _ := parseID(session.TenantID) - environment, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: tenant, ID: id}) - if err == nil { - value, err := sessionpg.EnvironmentFromRow(environment.Environment, environment.TenantID, environment.Configuration, nil) - if err != nil { - return session, err - } - session.Environment = &value - session.EnvironmentFailure = environmentFailure(environment.Environment) - state, err := sessionpg.LoadEnvironmentInput(ctx, q, id) - if err != nil { - return session, err - } - session.EnvironmentInputActivity, session.PendingInput = sessions.InputActivity(state) - } else if !errors.Is(err, pgx.ErrNoRows) { - return session, err - } - row, err := q.GetLatestSessionTurn(ctx, id) - if errors.Is(err, pgx.ErrNoRows) { - return session, nil - } - if err != nil { - return session, err - } - turn := sessionpg.TurnFromRow(row) - session.LastTurn = &turn - session.RequiredActions, err = sessionpg.LoadRequiredActions(ctx, q, id, turn) - if err != nil { - return session, err - } - session.Usage, err = q.SessionTokenUsage(ctx, id) - return session, err -} - -func environmentFailure(row sqlc.Environment) *sessions.EnvironmentFailure { - if row.Status != "failed" || !row.FailureReason.Valid || !row.FailedAt.Valid { - return nil - } - failure := &sessions.EnvironmentFailure{Reason: row.FailureReason.String, FailedAt: row.FailedAt.Time} - var detail sessions.ProvisioningFailureDetail - if json.Unmarshal(row.FailureDetail, &detail) == nil { - failure.Detail = sessions.SanitizedProvisioningDetail(detail) - } - return failure -} diff --git a/services/core/internal/store/self_hosted_cancel_public_test.go b/services/core/internal/store/self_hosted_cancel_public_test.go index eec981f6b..5da4bc866 100644 --- a/services/core/internal/store/self_hosted_cancel_public_test.go +++ b/services/core/internal/store/self_hosted_cancel_public_test.go @@ -36,7 +36,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { t.Helper() owner = executionOwner(t, db, s) worker, stop := publicOwnedWorker(t, s, db, owner) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(worker), executorURL("https://offline-executor.example")) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, worker), executorURL("https://offline-executor.example")) if err != nil { t.Fatal(err) } @@ -153,7 +153,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { if err != nil || len(itemsBefore.Items) != 2 { t.Fatal("controlled partial output was not recorded", err) } - cursor, err := s.SessionEventCursor(t.Context(), tenant, created.ID) + cursor, err := store.SessionAdapter(s).SessionEventCursor(t.Context(), tenant, created.ID) if err != nil { t.Fatal(err) } @@ -167,7 +167,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { if err != nil || !reflect.DeepEqual(itemsBefore, itemsAfter) { t.Fatal("cancellation admission changed partial history", err) } - afterCursor, err := s.SessionEventCursor(t.Context(), tenant, created.ID) + afterCursor, err := store.SessionAdapter(s).SessionEventCursor(t.Context(), tenant, created.ID) if err != nil || afterCursor != cursor { t.Fatal("cancellation admission fabricated an execution event", err) } diff --git a/services/core/internal/store/self_hosted_initial_public_test.go b/services/core/internal/store/self_hosted_initial_public_test.go index 90bcb379b..174f9d687 100644 --- a/services/core/internal/store/self_hosted_initial_public_test.go +++ b/services/core/internal/store/self_hosted_initial_public_test.go @@ -44,7 +44,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { t.Helper() enabled := []func(*api.Dependencies){acceptUnavailable(t)} if worker != nil { - enabled = append(enabled, workerExecution(worker), executorURL(origin)) + enabled = append(enabled, workerExecution(t, worker), executorURL(origin)) } else { // Without a Worker, Core keeps its executor URL but admits nothing. enabled = append(enabled, func(d *api.Dependencies) { diff --git a/services/core/internal/store/session_agent_filter_test.go b/services/core/internal/store/session_agent_filter_test.go index 959b6e2c1..b41aa7873 100644 --- a/services/core/internal/store/session_agent_filter_test.go +++ b/services/core/internal/store/session_agent_filter_test.go @@ -45,7 +45,7 @@ func TestSessionAgentFilterPaginationAndIsolation(t *testing.T) { var got []string cursor := "" for { - page, err := s.ListSessions(t.Context(), tenant, cursor, 2, ascending, &root) + page, err := sessionAdapter(s).ListSessions(t.Context(), tenant, cursor, 2, ascending, &root) if err != nil { t.Fatal(err) } @@ -75,20 +75,20 @@ func TestSessionAgentFilterPaginationAndIsolation(t *testing.T) { if got := read(s, false); !slices.Equal(got, reverse) { t.Fatal(got, reverse) } - unfiltered, err := s.ListSessions(t.Context(), tenant, "", 100, false, nil) + unfiltered, err := sessionAdapter(s).ListSessions(t.Context(), tenant, "", 100, false, nil) if err != nil || len(unfiltered.Sessions) != 9 { t.Fatal(unfiltered, err) } for _, id := range []string{"", "unknown", root + " ", "' OR true --"} { - page, err := s.ListSessions(t.Context(), tenant, "", 100, false, &id) + page, err := sessionAdapter(s).ListSessions(t.Context(), tenant, "", 100, false, &id) if err != nil || page.Sessions == nil || len(page.Sessions) != 0 || page.NextCursor != "" { t.Fatal(page, err) } } - if _, err := s.ListSessions(t.Context(), tenant, other.ID, 2, true, &root); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).ListSessions(t.Context(), tenant, other.ID, 2, true, &root); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign cursor", err) } - page, err := s.ListSessions(t.Context(), foreign, "", 100, false, &root) + page, err := sessionAdapter(s).ListSessions(t.Context(), foreign, "", 100, false, &root) if err != nil || len(page.Sessions) != 1 || page.Sessions[0].ID != other.ID { t.Fatal(page, err) } diff --git a/services/core/internal/store/session_artifacts_test.go b/services/core/internal/store/session_artifacts_test.go index 0c4fd6d8d..e35b125f7 100644 --- a/services/core/internal/store/session_artifacts_test.go +++ b/services/core/internal/store/session_artifacts_test.go @@ -176,7 +176,7 @@ func testSessionArtifactsPublishVersionScopeAndLifetime(t *testing.T, kind strin t.Fatalf("deleted metadata retained: %v", err) } } - if err := s.DeleteSession(t.Context(), tenant, session); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session}); err != nil { t.Fatal(err) } if count := largeObjectCount(t, pool); count != before { @@ -228,7 +228,7 @@ func TestSessionArtifactTransferDoesNotBlockDeletionOrCancellation(t *testing.T) want := sessions.ErrNotFound if operation == "delete" { // The idle-only decision itself is not blocked by the transfer. - if err := s.DeleteSession(ctx, tenant, session); !errors.Is(err, sessions.ErrNotIdle) { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatalf("transfer blocked or bypassed the deletion rule: %v", err) } if err := s.commitLegacyDeletion(ctx, tenant, session); err != nil { @@ -420,7 +420,7 @@ func TestSessionArtifactsRepublishOnlyNewChangedOrDeletedPaths(t *testing.T) { t.Fatalf("other Session first Turn published %v", got) } for _, id := range []string{session, other.ID} { - if err := s.DeleteSession(t.Context(), tenant, id); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: id}); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/session_configuration_test.go b/services/core/internal/store/session_configuration_test.go index d57de119b..b3047da67 100644 --- a/services/core/internal/store/session_configuration_test.go +++ b/services/core/internal/store/session_configuration_test.go @@ -22,11 +22,11 @@ func TestConfigurationSizeLimitSurvivesJSONBRoundTrip(t *testing.T) { if err != nil { t.Fatal(err) } - got, err := s.GetSession(ctx, tenant, first.ID) + got, err := sessionAdapter(s).GetSession(ctx, tenant, first.ID) if err != nil || string(got.Configuration) != string(first.Configuration) { t.Fatalf("configuration failed round trip: %v", err) } - page, err := s.ListSessions(ctx, tenant, "", 10, false, nil) + page, err := sessionAdapter(s).ListSessions(ctx, tenant, "", 10, false, nil) if err != nil || len(page.Sessions) != 1 || page.Sessions[0].ID != first.ID { t.Fatalf("configuration broke listing: %v", err) } @@ -61,7 +61,7 @@ func TestConfigurationIsPartOfSessionIdentity(t *testing.T) { t.Fatalf("changed snapshot was accepted: %v", err) } } - stored, err := s.GetSession(ctx, tenant, first.ID) + stored, err := sessionAdapter(s).GetSession(ctx, tenant, first.ID) if err != nil || string(stored.Configuration) != string(first.Configuration) { t.Fatalf("retry mutated snapshot: %+v, %v", stored, err) } diff --git a/services/core/internal/store/session_creation_stream_test.go b/services/core/internal/store/session_creation_stream_test.go index ce7b03c1d..fd0e70e19 100644 --- a/services/core/internal/store/session_creation_stream_test.go +++ b/services/core/internal/store/session_creation_stream_test.go @@ -53,7 +53,7 @@ func TestCreationStreamStartsBeforeOwnInputsAndRetriesAtUpsertCursor(t *testing. t.Fatal("invalid post-admission creation snapshot", created, retries) } id := created.Session.ID - initial, err := s.ListSessionEvents(ctx, tenant, id, created.Cursor) + initial, err := sessionAdapter(s).ListSessionEvents(ctx, tenant, id, created.Cursor) if err != nil || len(initial) != 3 { t.Fatal("lost initial events", initial, err) } @@ -72,7 +72,7 @@ func TestCreationStreamStartsBeforeOwnInputsAndRetriesAtUpsertCursor(t *testing. } return string(raw) } - read, err := s.GetSession(ctx, tenant, id) + read, err := sessionAdapter(s).GetSession(ctx, tenant, id) if err != nil { t.Fatal(err) } @@ -84,7 +84,7 @@ func TestCreationStreamStartsBeforeOwnInputsAndRetriesAtUpsertCursor(t *testing. if retry.Session.LastTurn != nil || retry.Session.EnvironmentInputActivity != nil || retry.Session.Usage != nil { t.Fatal("retry read the Session projection", retry.Session) } - if events, err := s.ListSessionEvents(ctx, tenant, id, retry.Cursor); err != nil || len(events) != 0 { + if events, err := sessionAdapter(s).ListSessionEvents(ctx, tenant, id, retry.Cursor); err != nil || len(events) != 0 { t.Fatal("retry replayed initial events", events, err) } } @@ -96,14 +96,14 @@ func TestCreationStreamStartsBeforeOwnInputsAndRetriesAtUpsertCursor(t *testing. if encoded(created.Session) != encoded(ordinary) || encoded(ordinary) != encoded(read) { t.Fatal("streamed and JSON creation projections differ", created.Session, ordinary) } - snapshot, cursor, err := s.SessionStreamSnapshot(ctx, tenant, id) + snapshot, cursor, err := sessionAdapter(s).SessionStreamSnapshot(ctx, tenant, id) if err != nil || encoded(snapshot) != encoded(read) || cursor != initial[len(initial)-1].Sequence || initial[2].Settled { t.Fatal("stream snapshot differs from the Session read and its cursor", cursor, err) } transition(t, s, tenant, id, ordinary.LastTurn.ID, sessions.TurnQueued, sessions.TurnInProgress) transition(t, s, tenant, id, ordinary.LastTurn.ID, sessions.TurnInProgress, sessions.TurnCompleted) // Completing before the HTTP observer drains does not change its start point. - all, err := s.ListSessionEvents(ctx, tenant, id, created.Cursor) + all, err := sessionAdapter(s).ListSessionEvents(ctx, tenant, id, created.Cursor) if err != nil || len(all) <= len(initial) || all[0].Event.EventID != initial[0].Event.EventID { t.Fatal(all, err) } @@ -111,10 +111,10 @@ func TestCreationStreamStartsBeforeOwnInputsAndRetriesAtUpsertCursor(t *testing. if last := all[len(all)-1]; last.Event.Type != "agent.session.idle" || !last.Settled { t.Fatal("terminal idle is not recorded as settled", last.Event.Type, last.Settled) } - if _, cursor, err := s.SessionStreamSnapshot(ctx, tenant, id); err != nil || cursor != all[len(all)-1].Sequence { + if _, cursor, err := sessionAdapter(s).SessionStreamSnapshot(ctx, tenant, id); err != nil || cursor != all[len(all)-1].Sequence { t.Fatal("stream snapshot cursor", cursor, err) } - if _, _, err := s.SessionStreamSnapshot(ctx, uuid.NewString(), id); !errors.Is(err, sessions.ErrNotFound) { + if _, _, err := sessionAdapter(s).SessionStreamSnapshot(ctx, uuid.NewString(), id); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign stream snapshot", err) } late, err := s.CreateSessionStream(ctx, tenant, input) @@ -128,7 +128,7 @@ func TestCreationStreamStartsBeforeOwnInputsAndRetriesAtUpsertCursor(t *testing. if late.Session.ID != id || late.Session.LastTurn != nil { t.Fatal("late retry read the Session projection", late.Session.LastTurn) } - future, err := s.ListSessionEvents(ctx, tenant, id, late.Cursor) + future, err := sessionAdapter(s).ListSessionEvents(ctx, tenant, id, late.Cursor) if err != nil || len(future) != 3 || future[0].Turn.ID != next[0].TurnID { t.Fatal(future, err) } diff --git a/services/core/internal/store/session_creator_test.go b/services/core/internal/store/session_creator_test.go index b380d1947..c8c0d25d6 100644 --- a/services/core/internal/store/session_creator_test.go +++ b/services/core/internal/store/session_creator_test.go @@ -23,7 +23,7 @@ func TestSessionCreatorIsRequiredBeforeCreation(t *testing.T) { t.Fatalf("invalid creator accepted: %v", err) } } - page, err := s.ListSessions(t.Context(), tenant, "", 10, false, nil) + page, err := sessionAdapter(s).ListSessions(t.Context(), tenant, "", 10, false, nil) if err != nil || len(page.Sessions) != 0 { t.Fatal("invalid creation wrote resources", page, err) } @@ -106,7 +106,7 @@ func TestConcurrentSessionCreatorsCannotShareCreationRetry(t *testing.T) { t.Fatal("early recovery ignored creator kind", err) } } - if _, err := restarted.GetSession(t.Context(), uuid.NewString(), winner.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(restarted).GetSession(t.Context(), uuid.NewString(), winner.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("creator bypassed project isolation", err) } } @@ -132,11 +132,11 @@ func TestHistoricalUnknownCreatorCannotBeClaimedByRetry(t *testing.T) { if err := pool.QueryRow(ctx, "SELECT to_jsonb(s)::text FROM sessions s WHERE id=$1", created.ID).Scan(&before); err != nil { t.Fatal(err) } - read, err := s.GetSession(ctx, tenant, created.ID) + read, err := sessionAdapter(s).GetSession(ctx, tenant, created.ID) if err != nil || read.Creator != nil { t.Fatal("historical ownership was invented", read, err) } - page, err := s.ListSessions(ctx, tenant, "", 10, false, nil) + page, err := sessionAdapter(s).ListSessions(ctx, tenant, "", 10, false, nil) if err != nil || len(page.Sessions) != 1 || page.Sessions[0].Creator != nil { t.Fatal("historical project reads changed", page, err) } diff --git a/services/core/internal/store/session_deletion.go b/services/core/internal/store/session_deletion.go deleted file mode 100644 index 7f6772abd..000000000 --- a/services/core/internal/store/session_deletion.go +++ /dev/null @@ -1,64 +0,0 @@ -package store - -import ( - "context" - "errors" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" - "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" -) - -// DeleteSession removes public access to a durably idle or failed Session while -// retaining state needed to settle execution. The decision is taken under the -// Session lock that also orders Turn and input admission, so a concurrent -// admission either commits first and is rejected here, or observes the deletion. -// The owner's repeated deletion succeeds without another resource write; foreign and -// missing Sessions remain not found. -func (s *Store) DeleteSession(ctx context.Context, tenantID, sessionID string) error { - return s.withLockedSession(ctx, tenantID, sessionID, true, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, locked sessions.LockedSession) error { - audit := func() error { - return auditpg.RecordWriteAudit(ctx, q, tenantID, "delete", "session", uuid.UUID(session.Bytes).String(), "") - } - if locked.Deleted { - return audit() - } - if err := requireSessionSettled(ctx, q, session); err != nil { - return err - } - if err := q.DeleteSessionArtifacts(ctx, session); err != nil { - return err - } - if err := q.ReleaseUnallocatedRuntimePlacement(ctx, session); err != nil { - return err - } - if err := q.MarkSessionDeleted(ctx, session); err != nil { - return err - } - return audit() - }) -} - -// requireSessionSettled rejects a queued, in-progress or waiting Turn, which -// includes pending required actions and function results, and a pending input -// reservation: queued later input, self-hosted input awaiting a connection, or -// hosted initial input while provisioning. Terminal idle and failed Sessions pass. -func requireSessionSettled(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - if _, err := q.GetActiveTurn(ctx, session); err == nil { - return sessions.ErrNotIdle - } else if !errors.Is(err, pgx.ErrNoRows) { - return err - } - state, err := sessionpg.LoadEnvironmentInput(ctx, q, session) - if err != nil { - return err - } - if _, pending := sessions.InputActivity(state); pending { - return sessions.ErrNotIdle - } - return nil -} diff --git a/services/core/internal/store/session_deletion_execution_test.go b/services/core/internal/store/session_deletion_execution_test.go index 678580d56..fe40df03a 100644 --- a/services/core/internal/store/session_deletion_execution_test.go +++ b/services/core/internal/store/session_deletion_execution_test.go @@ -31,7 +31,7 @@ func TestDeletedSessionWaitingTurnSettlesWithoutStoppingWorker(t *testing.T) { h.read(testExecutionRequest) h.write(input.TurnID, proto.TypeFunctionCall, proto.FunctionCallPayload{CallID: "pending", Name: "lookup_ticket", Arguments: json.RawMessage(`{}`)}) state := functionState(t, h, 1) - if err := h.s.DeleteSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := store.SessionService(t, h.s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("waiting Session deleted", err) } // A marker committed by an earlier release still cancels and settles work. @@ -52,7 +52,7 @@ func TestDeletedSessionWaitingTurnSettlesWithoutStoppingWorker(t *testing.T) { if err != nil || bound.NativeSessionID != "deleted-native" { t.Fatal(bound, err) } - if _, err := h.s.GetSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := store.SessionAdapter(h.s).GetSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } if _, err := h.bound().Run(ctx, h.tenant, h.session.ID, input.TurnID); !errors.Is(err, sessions.ErrNotFound) { @@ -72,7 +72,7 @@ func TestDeletedSessionRestartStillReconcilesHiddenClaim(t *testing.T) { if _, err := store.TransitionTurn(ctx, h.s, h.tenant, h.session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } - if err := h.s.DeleteSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := store.SessionService(t, h.s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("running Session deleted", err) } if err := h.s.CommitLegacyDeletion(ctx, h.tenant, h.session.ID); err != nil { @@ -88,7 +88,7 @@ func TestDeletedSessionRestartStillReconcilesHiddenClaim(t *testing.T) { if err != nil || turn.Status != sessions.TurnFailed { t.Fatal(turn, err) } - if _, err := h.s.GetSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := store.SessionAdapter(h.s).GetSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } } @@ -114,7 +114,7 @@ func TestWaitingSessionCancelsThenDeletesThroughWorker(t *testing.T) { h.read(testExecutionRequest) h.write(input.TurnID, proto.TypeFunctionCall, proto.FunctionCallPayload{CallID: "pending", Name: "lookup_ticket", Arguments: json.RawMessage(`{}`)}) state := functionState(t, h, 1) - if err := h.s.DeleteSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := store.SessionService(t, h.s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("waiting Session deleted", err) } if again := functionState(t, h, 1); again.LastTurn == nil || again.LastTurn.Status != sessions.TurnWaiting || !again.LastTurn.CancelRequestedAt.IsZero() { @@ -128,14 +128,14 @@ func TestWaitingSessionCancelsThenDeletesThroughWorker(t *testing.T) { if err := h.read(proto.TypePromptCancel).DecodePayload(&request); err != nil { t.Fatal(err) } - if err := h.s.DeleteSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := store.SessionService(t, h.s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("Session deleted before cancellation settled", err) } h.write(input.TurnID, proto.TypeInteractionDecisionAck, proto.InteractionDecisionAckPayload{DeliveryID: request.DeliveryID, Applied: true, Outcome: &proto.DonePayload{Metadata: map[string]any{proto.DoneMetaAgentSessionID: "cancelled-native"}}}) waitTurn(t, h, input.TurnID, sessions.TurnCancelled) deadline := time.Now().Add(10 * time.Second) for { - err := h.s.DeleteSession(ctx, h.tenant, h.session.ID) + err := store.SessionService(t, h.s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}) if err == nil { break } @@ -144,14 +144,14 @@ func TestWaitingSessionCancelsThenDeletesThroughWorker(t *testing.T) { } time.Sleep(20 * time.Millisecond) } - if err := h.s.DeleteSession(ctx, h.tenant, h.session.ID); err != nil { + if err := store.SessionService(t, h.s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}); err != nil { t.Fatal("repeated deletion", err) } raw, _ := json.Marshal(sessions.FunctionResultInput{TurnID: input.TurnID, CallID: state.RequiredActions[0].CallID, Result: json.RawMessage(`{"success":true,"output":"late"}`)}) if _, err := worker.SubmitInputs(ctx, h.tenant, h.session.ID, "late", []sessions.Input{{Kind: "tool_result", Payload: raw}}); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - if _, err := h.s.GetSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := store.SessionAdapter(h.s).GetSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } h.session = publicSession(t, h, "unrelated") diff --git a/services/core/internal/store/session_deletion_test.go b/services/core/internal/store/session_deletion_test.go index f1bd0de96..04b1b5623 100644 --- a/services/core/internal/store/session_deletion_test.go +++ b/services/core/internal/store/session_deletion_test.go @@ -77,23 +77,23 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { t.Fatal(err) } } - if err := s.DeleteSession(ctx, uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: uuid.NewString(), SessionID: session.ID}); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } if status == sessions.TurnQueued || status == sessions.TurnInProgress { // Deletion leaves active work untouched: no cancellation, marker or event. - cursor, err := s.SessionEventCursor(ctx, tenant, session.ID) + cursor, err := sessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } - if err := s.DeleteSession(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("active Session deleted", err) } turn, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, receipt.TurnID) if err != nil || turn.Status != status || !turn.CancelRequestedAt.IsZero() { t.Fatal("rejected deletion changed the Turn", turn, err) } - if after, err := s.SessionEventCursor(ctx, tenant, session.ID); err != nil || after != cursor { + if after, err := sessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID); err != nil || after != cursor { t.Fatal("rejected deletion recorded an event", after, cursor, err) } if sessionDeletedAt(t, pool, session.ID).Valid { @@ -105,7 +105,7 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { t.Fatal(err) } if status == sessions.TurnInProgress { - if err := s.DeleteSession(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("cancelling Session deleted", err) } if _, err := completeExecution(ctx, t, s, tenant, session.ID, receipt.TurnID, sessions.TurnCancelled, nil, "", receipt.Sequence); err != nil { @@ -113,24 +113,24 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { } } } - if err := s.DeleteSession(ctx, tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } marker := sessionDeletedAt(t, pool, session.ID) fresh := New(pool) // The owner's repeated deletion confirms again without another write. for _, repeat := range []*Store{s, fresh} { - if err := repeat.DeleteSession(ctx, tenant, session.ID); err != nil { + if err := sessionService(t, repeat).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal("repeated deletion", err) } - if err := repeat.DeleteSession(ctx, uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionService(t, repeat).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: uuid.NewString(), SessionID: session.ID}); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign deleted Session", err) } } if again := sessionDeletedAt(t, pool, session.ID); again != marker { t.Fatal("repeated deletion rewrote the marker", marker, again) } - if _, err := fresh.GetSession(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(fresh).GetSession(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } if _, err := fresh.CreateSession(ctx, tenant, input); !errors.Is(err, sessions.ErrIdempotencyConflict) { @@ -166,10 +166,10 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { if err != nil || len(inputs) == 0 || inputs[0].Sequence != receipt.Sequence { t.Fatal(inputs, err) } - if _, err := fresh.SessionEventCursor(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(fresh).SessionEventCursor(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - if _, err := fresh.ListSessionEvents(ctx, tenant, session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(fresh).ListSessionEvents(ctx, tenant, session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } }) @@ -290,7 +290,9 @@ func TestSessionDeletionRacesAdmissionUnderSessionLock(t *testing.T) { tenant, session := kind.setup(t, plain) deleted := make(chan error, 1) s := traced(t, func() { - go func() { deleted <- plain.DeleteSession(context.Background(), tenant, session) }() + go func() { + deleted <- sessionService(t, plain).DeleteSession(context.Background(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session}) + }() awaitSessionLockWaiter(t, pool) }) if err := kind.admit(t.Context(), s, tenant, session); err != nil { @@ -302,7 +304,7 @@ func TestSessionDeletionRacesAdmissionUnderSessionLock(t *testing.T) { if sessionDeletedAt(t, pool, session).Valid { t.Fatal("rejected deletion committed a marker") } - current, err := plain.GetSession(t.Context(), tenant, session) + current, err := sessionAdapter(plain).GetSession(t.Context(), tenant, session) if err != nil { t.Fatal(err) } @@ -320,7 +322,7 @@ func TestSessionDeletionRacesAdmissionUnderSessionLock(t *testing.T) { go func() { admitted <- kind.admit(context.Background(), plain, tenant, session) }() awaitSessionLockWaiter(t, pool) }) - if err := s.DeleteSession(t.Context(), tenant, session); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session}); err != nil { t.Fatal(err) } if err := <-admitted; !errors.Is(err, sessions.ErrNotFound) { @@ -341,7 +343,10 @@ func TestSessionDeletionRacesAdmissionUnderSessionLock(t *testing.T) { other := New(pool) start := make(chan struct{}) results := make(chan error, 2) - go func() { <-start; results <- plain.DeleteSession(context.Background(), tenant, session) }() + go func() { + <-start + results <- sessionService(t, plain).DeleteSession(context.Background(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session}) + }() go func() { <-start; results <- kind.admit(context.Background(), other, tenant, session) }() close(start) first, second := <-results, <-results @@ -406,10 +411,10 @@ func TestSessionDeletionKeepsProvisioningInputPlacementUntilSettled(t *testing.T if before.deleted.Valid || before.released.Valid || before.retained != 1 || before.reserved != 1 { t.Fatal("unexpected reserved placement", before) } - if err := s.DeleteSession(ctx, uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: uuid.NewString(), SessionID: session.ID}); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign deletion", err) } - if err := s.DeleteSession(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("provisioning input deleted", err) } if after := read(); after != before { @@ -421,20 +426,20 @@ func TestSessionDeletionKeepsProvisioningInputPlacementUntilSettled(t *testing.T if count, err := w.ExpireEnvironmentInputs(ctx); err != nil || count != 1 { t.Fatal("initial input did not expire", count, err) } - if err := s.DeleteSession(ctx, tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } deleted := read() if !deleted.deleted.Valid || !deleted.released.Valid || deleted.retained != 0 || deleted.reserved != 0 { t.Fatal("allowed deletion kept the placement", deleted) } - if err := s.DeleteSession(ctx, tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal("repeated deletion", err) } if again := read(); again != deleted { t.Fatal("repeated deletion changed timestamps", deleted, again) } - if err := s.DeleteSession(ctx, uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: uuid.NewString(), SessionID: session.ID}); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign deletion of a deleted Session", err) } } diff --git a/services/core/internal/store/session_diagnostics.go b/services/core/internal/store/session_diagnostics.go deleted file mode 100644 index 0c9ee313a..000000000 --- a/services/core/internal/store/session_diagnostics.go +++ /dev/null @@ -1,87 +0,0 @@ -package store - -import ( - "context" - "errors" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" - "github.com/jackc/pgx/v5" -) - -// GetSessionDiagnosticsSnapshot keeps all existing Session projections in one -// read-only snapshot, including the failure precedence used by the public API. -func (s *Store) GetSessionDiagnosticsSnapshot(ctx context.Context, tenantID, sessionID string) (sessions.Session, error) { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.Session{}, err - } - var session sessions.Session - err = s.pooled.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - row, err := q.GetSession(ctx, sqlc.GetSessionParams{TenantID: tenant, ID: pgunit.PathID(sessionID)}) - if err != nil { - return err - } - session, err = sessionpg.SessionFromRow(row) - if err != nil { - return err - } - session, err = readSessionActivity(ctx, q, session) - return err - }) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.Session{}, sessions.ErrNotFound - } - return session, err -} - -// GetTurnDiagnosticsSnapshot reads only root Turns and their root Items. Its -// bounded query and public projection share the same committed snapshot. -func (s *Store) GetTurnDiagnosticsSnapshot(ctx context.Context, tenantID, sessionID, turnID string) (sessions.TurnDiagnosticsSnapshot, error) { - params, err := sessionpg.PublicTurnLookup(tenantID, sessionID, turnID) - if err != nil { - return sessions.TurnDiagnosticsSnapshot{}, err - } - result := sessions.TurnDiagnosticsSnapshot{Items: []sessions.ItemDiagnosticTiming{}} - err = s.pooled.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - turn, err := q.GetTurn(ctx, params) - if err != nil { - return err - } - row, err := q.GetSession(ctx, sqlc.GetSessionParams{TenantID: params.TenantID, ID: params.SessionID}) - if err != nil { - return err - } - result.Session, err = sessionpg.SessionFromRow(row) - if err != nil { - return err - } - result.Turn = sessionpg.TurnFromRow(turn) - rows, err := q.ListTurnItemDiagnostics(ctx, sqlc.ListTurnItemDiagnosticsParams{SessionID: params.SessionID, TurnID: params.ID}) - if err != nil { - return err - } - result.ItemsTruncated = len(rows) > 1000 - if result.ItemsTruncated { - rows = rows[:1000] - } - for _, row := range rows { - item := sessions.ItemDiagnosticTiming{ItemID: uuid.UUID(row.ID.Bytes).String(), StartedAt: row.CreatedAt.Time} - if row.SettledAt.Valid { - at := row.SettledAt.Time - item.CompletedAt = &at - } - result.Items = append(result.Items, item) - } - return nil - }) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.TurnDiagnosticsSnapshot{}, sessions.ErrNotFound - } - return result, err -} diff --git a/services/core/internal/store/session_diagnostics_cancellation_test.go b/services/core/internal/store/session_diagnostics_cancellation_test.go index f4aae5f19..03518915a 100644 --- a/services/core/internal/store/session_diagnostics_cancellation_test.go +++ b/services/core/internal/store/session_diagnostics_cancellation_test.go @@ -33,9 +33,9 @@ func TestDiagnosticSnapshotCancellationReleasesTransaction(t *testing.T) { go func() { var err error if rootTurn { - _, err = s.GetTurnDiagnosticsSnapshot(ctx, tenant, session.ID, receipt.TurnID) + _, err = sessionAdapter(s).GetTurnDiagnosticsSnapshot(ctx, tenant, session.ID, receipt.TurnID) } else { - _, err = s.GetSessionDiagnosticsSnapshot(ctx, tenant, session.ID) + _, err = sessionAdapter(s).GetSession(ctx, tenant, session.ID) } done <- err }() @@ -81,9 +81,9 @@ func TestDiagnosticSnapshotCancellationReleasesTransaction(t *testing.T) { t.Fatal(err) } if rootTurn { - _, err = s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + _, err = sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) } else { - _, err = s.GetSessionDiagnosticsSnapshot(t.Context(), tenant, session.ID) + _, err = sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) } if err != nil { t.Fatal("snapshot read did not recover", err) diff --git a/services/core/internal/store/session_diagnostics_test.go b/services/core/internal/store/session_diagnostics_test.go index 9757c2dfe..395f8719c 100644 --- a/services/core/internal/store/session_diagnostics_test.go +++ b/services/core/internal/store/session_diagnostics_test.go @@ -34,7 +34,7 @@ func TestDiagnosticItemReceiptSettlementAndReplay(t *testing.T) { t.Fatal(err) } } - snap, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + snap, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil || len(snap.Items) != 2 { t.Fatal(snap, err) } @@ -58,7 +58,7 @@ func TestDiagnosticItemReceiptSettlementAndReplay(t *testing.T) { t.Fatal(err) } transition(t, s, tenant, session.ID, receipt.TurnID, sessions.TurnInProgress, sessions.TurnFailed) - snap, err = s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + snap, err = sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil || snap.Items[1].CompletedAt != nil { t.Fatal("historical settlement synthesized", snap, err) } @@ -81,7 +81,7 @@ func TestDiagnosticForceSettlementIgnoresNativeClock(t *testing.T) { if err != nil || completed.CompletedAt.UnixMilli() != source { t.Fatal("public native completion changed", completed, err) } - snap, err := s.GetTurnDiagnosticsSnapshot(t.Context(), owner.TenantID, owner.SessionID, turn) + snap, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), owner.TenantID, owner.SessionID, turn) if err != nil || len(snap.Items) != 2 { t.Fatal(snap, err) } @@ -145,7 +145,7 @@ func TestDiagnosticSettlementWaitsForSessionLock(t *testing.T) { if err = <-done; err != nil { t.Fatal(err) } - snap, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + snap, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil { t.Fatal(err) } @@ -165,7 +165,7 @@ func TestDiagnosticTimingBoundOrderAndIsolation(t *testing.T) { if err != nil { t.Fatal(err) } - snap, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + snap, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil || len(snap.Items) != 1000 || !snap.ItemsTruncated { t.Fatal("unbounded diagnostics", len(snap.Items), snap.ItemsTruncated, err) } @@ -179,20 +179,20 @@ func TestDiagnosticTimingBoundOrderAndIsolation(t *testing.T) { } } runtimeSuspensionSQL(t, pool, "DELETE FROM session_items WHERE turn_id=$1 AND created_at>'2020-01-01T00:00:00Z'", receipt.TurnID) - exact, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + exact, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil || len(exact.Items) != 1000 || exact.ItemsTruncated { t.Fatal("exact limit falsely truncated", len(exact.Items), exact.ItemsTruncated, err) } for _, ids := range [][3]string{{uuid.NewString(), session.ID, receipt.TurnID}, {tenant, "malformed", receipt.TurnID}, {tenant, session.ID, uuid.NewString()}} { - if _, err := s.GetTurnDiagnosticsSnapshot(t.Context(), ids[0], ids[1], ids[2]); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), ids[0], ids[1], ids[2]); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("scope leaked", ids, err) } } runtimeSuspensionSQL(t, pool, "UPDATE sessions SET deleted_at=clock_timestamp() WHERE id=$1", session.ID) - if _, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted root visible", err) } - if _, err := s.GetSessionDiagnosticsSnapshot(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted Session visible", err) } } @@ -233,11 +233,11 @@ func TestDiagnosticProvisioningDetailAtomicAndPrivate(t *testing.T) { if err = sessionExecution(t, writer.lease).FailEnvironmentInitialization(t.Context(), preparation, failure); err != nil { t.Fatal(err) } - snap, err := s.GetSessionDiagnosticsSnapshot(t.Context(), tenant, session.ID) + snap, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || snap.EnvironmentFailure == nil || snap.EnvironmentFailure.Detail == nil || *snap.EnvironmentFailure.Detail.Index != 2 || *snap.EnvironmentFailure.Detail.ExitCode != 7 { t.Fatal("detail not persisted", snap.EnvironmentFailure, err) } - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + events, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil { t.Fatal(err) } @@ -247,7 +247,7 @@ func TestDiagnosticProvisioningDetailAtomicAndPrivate(t *testing.T) { } // Historical reasons are never parsed into structured detail; corrupted private fields are sanitized. runtimeSuspensionSQL(t, pool, "UPDATE environments SET failure_detail=$2 WHERE id=$1", owner.EnvironmentID, `{"step":"secret-provider-token","index":3,"exit_code":2}`) - snap, err = s.GetSessionDiagnosticsSnapshot(t.Context(), tenant, session.ID) + snap, err = sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil || snap.EnvironmentFailure.Detail != nil { t.Fatal("unsafe private detail projected", snap.EnvironmentFailure, err) } @@ -271,7 +271,7 @@ func TestDiagnosticPutItemRollsBack(t *testing.T) { if err = tx.Rollback(t.Context()); err != nil { t.Fatal(err) } - snap, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + snap, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil || len(snap.Items) != 1 { t.Fatal("rolled back receipt visible", snap, err) } @@ -287,7 +287,7 @@ func TestDiagnosticFirstSettlementSurvivesStoredStatusRegression(t *testing.T) { if err := journal.AppendTurnEvents(t.Context(), tenant, session.ID, receipt.TurnID, 1, events); err != nil { t.Fatal(err) } - first, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + first, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil { t.Fatal(err) } @@ -304,7 +304,7 @@ func TestDiagnosticFirstSettlementSurvivesStoredStatusRegression(t *testing.T) { } else if err := journal.AppendTurnEvents(t.Context(), tenant, session.ID, receipt.TurnID, 3, events[1:]); err != nil { t.Fatal(err) } - got, err := s.GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) + got, err := sessionAdapter(s).GetTurnDiagnosticsSnapshot(t.Context(), tenant, session.ID, receipt.TurnID) if err != nil || got.Items[1].CompletedAt == nil || !got.Items[1].CompletedAt.Equal(*item.CompletedAt) { t.Fatal("first settlement overwritten", force, got, err) } @@ -318,7 +318,7 @@ func TestDiagnosticRootReadRejectsActualChildTurn(t *testing.T) { runtimeSuspensionSQL(t, s.pool, `INSERT INTO turn_events(session_id,turn_id,ordinal,kind,payload) VALUES($1,$2,1,'subagent','{}')`, owner.SessionID, root) runtimeSuspensionSQL(t, s.pool, `INSERT INTO subagent_identities(id,session_id,device_id,engine,native_id,parent_native_id,native_created_at,first_turn_id,first_event_ordinal) VALUES($1,$2,$3,'codex','child','root',1,$4,1)`, child, owner.SessionID, owner.DeviceID, root) runtimeSuspensionSQL(t, s.pool, `INSERT INTO subagent_turns(id,session_id,subagent_id,native_id,status,created_at) VALUES($1,$2,$3,'child-turn','in_progress',clock_timestamp())`, turn, owner.SessionID, child) - if _, err := w.GetTurnDiagnosticsSnapshot(t.Context(), owner.TenantID, owner.SessionID, turn); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(w).GetTurnDiagnosticsSnapshot(t.Context(), owner.TenantID, owner.SessionID, turn); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("child Turn became root diagnostics", err) } } @@ -357,7 +357,7 @@ func TestDiagnosticSessionSnapshotConcurrentCommitConsistency(t *testing.T) { done <- nil }() for i := 0; i < 150; i++ { - snapshot, err := s.GetSessionDiagnosticsSnapshot(t.Context(), tenant, session.ID) + snapshot, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_environment_snapshot_test.go b/services/core/internal/store/session_environment_snapshot_test.go index 39d54f2d2..56e9ce389 100644 --- a/services/core/internal/store/session_environment_snapshot_test.go +++ b/services/core/internal/store/session_environment_snapshot_test.go @@ -29,7 +29,7 @@ func TestSelfHostedCreationSnapshotRetainsEnvironmentAndCursor(t *testing.T) { if err != nil || pending.State != sessions.EnvironmentInputPending { t.Fatal(pending, err) } - events, err := s.ListSessionEvents(t.Context(), tenant, created.Session.ID, created.Cursor) + events, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, created.Session.ID, created.Cursor) if err != nil || len(events) != 1 || events[0].Event.Type != "agent.session.requires_action" { t.Fatal("creation cursor lost subsequent activity", events, err) } @@ -59,7 +59,7 @@ func TestSelfHostedCreationSnapshotRetainsEnvironmentAndCursor(t *testing.T) { if _, err := s.FindSessionCreation(t.Context(), tenant, input.IdempotencyKey, input.CreationRequest, changedCreator); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal("retry creator isolation", err) } - current, err := s.GetSession(t.Context(), tenant, created.Session.ID) + current, err := sessionAdapter(s).GetSession(t.Context(), tenant, created.Session.ID) if err != nil || current.EnvironmentInputActivity == nil || current.EnvironmentInputActivity.Status != "requires_action" { t.Fatal("ordinary read lost current activity", current, err) } diff --git a/services/core/internal/store/session_events.go b/services/core/internal/store/session_events.go deleted file mode 100644 index 81c48d111..000000000 --- a/services/core/internal/store/session_events.go +++ /dev/null @@ -1,67 +0,0 @@ -package store - -import ( - "bytes" - "context" - "encoding/json" - "errors" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" -) - -func (s *Store) SessionEventCursor(ctx context.Context, tenantID, sessionID string) (int64, error) { - tenant, err := parseID(tenantID) - if err != nil { - return 0, err - } - id := pgunit.PathID(sessionID) - cursor, err := s.queries.SessionEventCursor(ctx, sqlc.SessionEventCursorParams{TenantID: tenant, ID: id}) - if errors.Is(err, pgx.ErrNoRows) { - return 0, sessions.ErrNotFound - } - return cursor, err -} - -func (s *Store) ListSessionEvents(ctx context.Context, tenantID, sessionID string, after int64) ([]sessions.SessionChange, error) { - if after < 0 { - return nil, sessions.ErrInvalidInput - } - latest, err := s.SessionEventCursor(ctx, tenantID, sessionID) - if err != nil { - return nil, err - } - tenant, err := parseID(tenantID) - if err != nil { - return nil, err - } - id, err := parseID(sessionID) - if err != nil { - return nil, err - } - rows, err := s.queries.ListSessionEvents(ctx, sqlc.ListSessionEventsParams{TenantID: tenant, SessionID: id, Sequence: after}) - if err != nil { - return nil, err - } - changes := make([]sessions.SessionChange, 0, len(rows)) - if len(rows) == 0 && latest > after { - return nil, sessions.ErrStreamGap - } - for _, row := range rows { - if row.Sequence != after+1 { - return nil, sessions.ErrStreamGap - } - var change sessions.SessionChange - decoder := json.NewDecoder(bytes.NewReader(row.Payload)) - decoder.UseNumber() - if err := decoder.Decode(&change); err != nil { - return nil, err - } - change.Sequence = row.Sequence - changes = append(changes, change) - after = row.Sequence - } - return changes, nil -} diff --git a/services/core/internal/store/session_events_test.go b/services/core/internal/store/session_events_test.go index 48fa34581..3b7748f58 100644 --- a/services/core/internal/store/session_events_test.go +++ b/services/core/internal/store/session_events_test.go @@ -31,7 +31,7 @@ func TestSessionEventsAreVisibleOnlyAfterCommit(t *testing.T) { if err = sqlc.New(tx).AppendSessionEvent(ctx, sqlc.AppendSessionEventParams{ID: pgtype.UUID{Bytes: uuid.MustParse(session.ID), Valid: true}, Payload: []byte(`{"event":{"type":"agent.session.idle"}}`)}); err != nil { t.Fatal(err) } - if changes, err := s.ListSessionEvents(ctx, tenant, session.ID, 0); err != nil || len(changes) != 0 { + if changes, err := store.SessionAdapter(s).ListSessionEvents(ctx, tenant, session.ID, 0); err != nil || len(changes) != 0 { t.Fatal("uncommitted events were visible", changes, err) } if commit { @@ -43,7 +43,7 @@ func TestSessionEventsAreVisibleOnlyAfterCommit(t *testing.T) { t.Fatal(err) } } - if changes, err := s.ListSessionEvents(ctx, tenant, session.ID, 0); err != nil || len(changes) != 1 || changes[0].Sequence != 1 { + if changes, err := store.SessionAdapter(s).ListSessionEvents(ctx, tenant, session.ID, 0); err != nil || len(changes) != 1 || changes[0].Sequence != 1 { t.Fatal("commit or rollback changed sequence continuity", changes, err) } } @@ -61,14 +61,14 @@ func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { if err != nil { t.Fatal(err) } - before, err := s.SessionEventCursor(ctx, tenant, session.ID) + before, err := store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } if _, err = s.SubmitMessage(ctx, tenant, session.ID, "start", json.RawMessage(`{"text":"question"}`)); err != nil { t.Fatal(err) } - after, _ := s.SessionEventCursor(ctx, tenant, session.ID) + after, _ := store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if before != after { t.Fatal("input retry published duplicate events") } @@ -86,14 +86,14 @@ func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { t.Fatal(err) } } - before, _ = s.SessionEventCursor(ctx, tenant, session.ID) + before, _ = store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if err = journal.AppendTurnEvents(ctx, tenant, session.ID, input.TurnID, 5, []sessions.ExecutionEvent{ {Kind: "delta", Payload: json.RawMessage(`{"item_id":"discarded","delta":"rollback"}`)}, {Kind: "output_message", Payload: json.RawMessage(`{"status":"invalid"}`)}, }); err == nil { t.Fatal("invalid projection accepted") } - after, _ = s.SessionEventCursor(ctx, tenant, session.ID) + after, _ = store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if before != after { t.Fatal("failed transaction published events") } @@ -103,7 +103,7 @@ func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { var all []sessions.SessionChange cursor := int64(0) for { - page, err := store.New(pool).ListSessionEvents(ctx, tenant, session.ID, cursor) + page, err := store.SessionAdapter(store.New(pool)).ListSessionEvents(ctx, tenant, session.ID, cursor) if err != nil { t.Fatal(err) } @@ -139,14 +139,14 @@ func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { if counts["agent.session.turn.output_text.delta"] != 3 || counts["agent.session.turn.item.done"] != 2 { t.Fatal(counts) } - if _, err = s.ListSessionEvents(ctx, uuid.NewString(), session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { + if _, err = store.SessionAdapter(s).ListSessionEvents(ctx, uuid.NewString(), session.ID, 0); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign event access", err) } - before, _ = s.SessionEventCursor(ctx, tenant, session.ID) + before, _ = store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if _, err = sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true); err != nil { t.Fatal(err) } - after, _ = s.SessionEventCursor(ctx, tenant, session.ID) + after, _ = store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if before != after { t.Fatal("history read published live events") } @@ -178,17 +178,17 @@ func TestSessionEventsRetentionAndQueuedCancellation(t *testing.T) { if err = pool.QueryRow(ctx, "SELECT count(*) FROM session_events WHERE session_id=$1", session.ID).Scan(&count); err != nil || count != 256 { t.Fatal(count, err) } - if _, err = s.ListSessionEvents(ctx, tenant, session.ID, 0); !errors.Is(err, sessions.ErrStreamGap) { + if _, err = store.SessionAdapter(s).ListSessionEvents(ctx, tenant, session.ID, 0); !errors.Is(err, sessions.ErrStreamGap) { t.Fatal("lagging reader did not detect missing events", err) } - cursor, err := s.SessionEventCursor(ctx, tenant, session.ID) + cursor, err := store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } if _, err = s.RequestCancel(ctx, tenant, session.ID, "cancel"); err != nil { t.Fatal(err) } - changes, err := s.ListSessionEvents(ctx, tenant, session.ID, cursor) + changes, err := store.SessionAdapter(s).ListSessionEvents(ctx, tenant, session.ID, cursor) if err != nil || len(changes) != 2 || changes[0].Event.Type != "agent.session.turn.cancelled" || changes[1].Event.Type != "agent.session.idle" { t.Fatal(changes, err) } diff --git a/services/core/internal/store/session_execution_configuration.go b/services/core/internal/store/session_execution_configuration.go index da0cc42fc..883c44bf6 100644 --- a/services/core/internal/store/session_execution_configuration.go +++ b/services/core/internal/store/session_execution_configuration.go @@ -3,15 +3,12 @@ package store import ( "context" "encoding/json" - "errors" "fmt" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" ) @@ -23,7 +20,7 @@ func saveSessionExecutionConfiguration(ctx context.Context, q *sqlc.Queries, ses } var frozenRevision pgtype.UUID frozen := *projection - model, err := sessionExecutionModel(session.Configuration) + model, err := sessions.ExecutionModel(session.Configuration) if err != nil { return err } @@ -54,7 +51,7 @@ func saveSessionExecutionConfiguration(ctx context.Context, q *sqlc.Queries, ses default: return fmt.Errorf("%w: invalid execution projection source", sessions.ErrInvalidInput) } - normalizeExecutionProjection(&frozen, uuid.UUID(session.ID.Bytes).String()) + sessions.NormalizeExecutionProjection(&frozen, uuid.UUID(session.ID.Bytes).String()) raw, err := json.Marshal(frozen) if err != nil { return err @@ -62,69 +59,6 @@ func saveSessionExecutionConfiguration(ctx context.Context, q *sqlc.Queries, ses return q.SaveSessionExecutionConfiguration(ctx, sqlc.SaveSessionExecutionConfigurationParams{SessionID: session.ID, Configuration: raw, DeploymentProviderRevision: frozenRevision}) } -// GetSessionExecutionConfiguration reads only safe committed configuration. It -// never loads provider ciphertext, current defaults or runtime health. -func (s *Store) GetSessionExecutionConfiguration(ctx context.Context, tenantID, sessionID string) (v1.SessionExecutionConfiguration, error) { - tenant, err := parseID(tenantID) - if err != nil { - return v1.SessionExecutionConfiguration{}, err - } - row, err := s.queries.GetSessionExecutionConfiguration(ctx, sqlc.GetSessionExecutionConfigurationParams{TenantID: tenant, SessionID: pgunit.PathID(sessionID)}) - if errors.Is(err, pgx.ErrNoRows) { - return v1.SessionExecutionConfiguration{}, sessions.ErrNotFound - } - if err != nil { - return v1.SessionExecutionConfiguration{}, fmt.Errorf("get session execution configuration: %w", err) - } - var projection v1.SessionExecutionConfiguration - if len(row.ExecutionConfiguration) == 0 { - model, err := sessionExecutionModel(row.SessionConfiguration) - if err != nil { - return projection, err - } - var harness *string - if row.Engine != "" { - harness = &row.Engine - } - projection.Model = v1.ExecutionSelection{Value: model, Source: "unknown"} - projection.Harness = v1.ExecutionSelection{Value: harness, Source: "unknown"} - projection.ModelProvider = v1.ExecutionProviderSelection{Source: "unknown", Status: "unavailable"} - } else if err := json.Unmarshal(row.ExecutionConfiguration, &projection); err != nil { - return v1.SessionExecutionConfiguration{}, errors.New("invalid stored session execution configuration") - } - normalizeExecutionProjection(&projection, uuid.UUID(row.ID.Bytes).String()) - return projection, nil -} - -func normalizeExecutionProjection(projection *v1.SessionExecutionConfiguration, sessionID string) { - projection.Object = "agent.session.execution_configuration" - projection.SchemaVersion = 1 - if projection.HarnessConfig.Source == "" { - projection.HarnessConfig.Source = "unknown" - } - projection.HarnessConfig.Value = v1.ResolvedHarnessConfig(projection.HarnessConfig.Value) - projection.SessionID = sessionID - if projection.ModelProvider.Source == "deployment" && (projection.ModelProvider.Status != "available" || projection.ModelProvider.Configuration == nil) { - // Sessions created before deployment defaults moved into Core stay redacted. - projection.ModelProvider.Status = "redacted" - projection.ModelProvider.Configuration = nil - } else if projection.ModelProvider.Status != "available" { - projection.ModelProvider.Configuration = nil - } -} - -func sessionExecutionModel(configuration []byte) (*string, error) { - var config struct { - Agent struct { - Model *string `json:"model"` - } `json:"agent"` - } - if err := json.Unmarshal(configuration, &config); err != nil { - return nil, errors.New("invalid stored Session model configuration") - } - return config.Agent.Model, nil -} - func sameExecutionValue(a, b *string) bool { return a == nil && b == nil || a != nil && b != nil && *a == *b } diff --git a/services/core/internal/store/session_execution_configuration_test.go b/services/core/internal/store/session_execution_configuration_test.go index e15ef2f49..0c9338d1e 100644 --- a/services/core/internal/store/session_execution_configuration_test.go +++ b/services/core/internal/store/session_execution_configuration_test.go @@ -61,7 +61,7 @@ func TestSessionExecutionConfigurationFrozenAcrossCreationPathsAndRetry(t *testi } // A reader without the encryption key can use the safe snapshot after restart. reader := New(pool) - frozen, err := reader.GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) + frozen, err := sessionAdapter(reader).GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -90,13 +90,13 @@ func TestSessionExecutionConfigurationFrozenAcrossCreationPathsAndRetry(t *testi if err != nil || replay.ID != session.ID { t.Fatal("projection metadata changed retry identity", err) } - if _, err := s.UpdateSessionMetadata(t.Context(), tenant, session.ID, map[string]string{"edited": "yes"}); err != nil { + if _, err := sessionService(t, s).UpdateSessionMetadata(t.Context(), sessions.UpdateSessionMetadataCommand{TenantID: tenant, SessionID: session.ID, Metadata: map[string]string{"edited": "yes"}}); err != nil { t.Fatal(err) } if _, err := pool.Exec(t.Context(), "UPDATE session_model_execution SET encrypted_config='\\x00'::bytea WHERE session_id=$1", session.ID); err != nil { t.Fatal(err) } - got, err := reader.GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) + got, err := sessionAdapter(reader).GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) if err != nil || !reflect.DeepEqual(got, frozen) { t.Fatal("snapshot changed or reader decrypted a secret", err) } @@ -105,19 +105,19 @@ func TestSessionExecutionConfigurationFrozenAcrossCreationPathsAndRetry(t *testi if _, err := pool.Exec(t.Context(), `UPDATE session_execution_configuration SET configuration = jsonb_set(configuration, '{model_provider}', '{"source":"deployment","status":"redacted","configuration":null}') WHERE session_id=$1`, session.ID); err != nil { t.Fatal(err) } - if historical, err := reader.GetSessionExecutionConfiguration(t.Context(), tenant, session.ID); err != nil || historical.ModelProvider.Status != "redacted" || historical.ModelProvider.Configuration != nil { + if historical, err := sessionAdapter(reader).GetSessionExecutionConfiguration(t.Context(), tenant, session.ID); err != nil || historical.ModelProvider.Status != "redacted" || historical.ModelProvider.Configuration != nil { t.Fatal("historical deployment selection changed", err) } } for _, lookup := range []struct{ tenant, id string }{{uuid.NewString(), session.ID}, {tenant, uuid.NewString()}, {tenant, "malformed"}} { - if _, err := reader.GetSessionExecutionConfiguration(t.Context(), lookup.tenant, lookup.id); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(reader).GetSessionExecutionConfiguration(t.Context(), lookup.tenant, lookup.id); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("foreign/missing projection read differed", err) } } - if err := s.DeleteSession(t.Context(), tenant, session.ID); err != nil { + if err := sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { t.Fatal(err) } - if _, err := reader.GetSessionExecutionConfiguration(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(reader).GetSessionExecutionConfiguration(t.Context(), tenant, session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted Session projection remained public", err) } }) @@ -134,7 +134,7 @@ func TestSessionExecutionConfigurationHistoricalProvenance(t *testing.T) { if err != nil { t.Fatal(err) } - got, err := s.GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) + got, err := sessionAdapter(s).GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -222,7 +222,7 @@ func TestSessionExecutionConfigurationConcurrentRetryKeepsWinner(t *testing.T) { t.Error(err) return } - got, err := s.GetSessionExecutionConfiguration(t.Context(), tenant, created.Session.ID) + got, err := sessionAdapter(s).GetSessionExecutionConfiguration(t.Context(), tenant, created.Session.ID) if err != nil { t.Error(err) return @@ -253,7 +253,7 @@ func TestSessionExecutionConfigurationSurvivesSuspendResume(t *testing.T) { if err != nil { t.Fatal(err) } - frozen, err := s.GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) + frozen, err := sessionAdapter(s).GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) if err != nil { t.Fatal(err) } @@ -285,7 +285,7 @@ func TestSessionExecutionConfigurationSurvivesSuspendResume(t *testing.T) { if err != nil { t.Fatal(err) } - got, err := s.GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) + got, err := sessionAdapter(s).GetSessionExecutionConfiguration(t.Context(), tenant, session.ID) if err != nil || !reflect.DeepEqual(got, frozen) { t.Fatal("runtime transition changed execution projection", phase, err) } diff --git a/services/core/internal/store/session_initial_input_test.go b/services/core/internal/store/session_initial_input_test.go index e581f228d..040701873 100644 --- a/services/core/internal/store/session_initial_input_test.go +++ b/services/core/internal/store/session_initial_input_test.go @@ -74,10 +74,10 @@ func TestInitialInputCreationRetriesAcrossConnectionsAndLaterTurns(t *testing.T) if _, err := s.RequestCancel(ctx, tenant, first.ID, "cancel"); err != nil { t.Fatal(err) } - if _, err := s.UpdateSessionMetadata(ctx, tenant, first.ID, map[string]string{"updated": "yes"}); err != nil { + if _, err := sessionService(t, s).UpdateSessionMetadata(ctx, sessions.UpdateSessionMetadataCommand{TenantID: tenant, SessionID: first.ID, Metadata: map[string]string{"updated": "yes"}}); err != nil { t.Fatal(err) } - events, err := s.ListSessionEvents(ctx, tenant, first.ID, 0) + events, err := sessionAdapter(s).ListSessionEvents(ctx, tenant, first.ID, 0) if err != nil { t.Fatal(err) } @@ -87,7 +87,7 @@ func TestInitialInputCreationRetriesAcrossConnectionsAndLaterTurns(t *testing.T) if err != nil || retry.ID != first.ID || retry.LastTurn.ID != next[0].TurnID || retry.LastTurn.Status != sessions.TurnCancelled || retry.Metadata["updated"] != "yes" { t.Fatal(retry, err) } - after, err := restarted.ListSessionEvents(ctx, tenant, first.ID, 0) + after, err := sessionAdapter(restarted).ListSessionEvents(ctx, tenant, first.ID, 0) if err != nil || !reflect.DeepEqual(events, after) { t.Fatal("retry emitted more events", err) } @@ -112,7 +112,7 @@ func TestInitialInputFailureRollsBackSessionAndWork(t *testing.T) { if got, err := s.CreateSession(ctx, tenant, input); err == nil || got.ID != "" { t.Fatal("partial creation succeeded", got, err) } - page, err := s.ListSessions(ctx, tenant, "", 100, true, nil) + page, err := sessionAdapter(s).ListSessions(ctx, tenant, "", 100, true, nil) if err != nil || len(page.Sessions) != 0 { t.Fatal("partial Session survived", page, err) } diff --git a/services/core/internal/store/session_initial_public_test.go b/services/core/internal/store/session_initial_public_test.go index da77de17e..fc2bb5622 100644 --- a/services/core/internal/store/session_initial_public_test.go +++ b/services/core/internal/store/session_initial_public_test.go @@ -29,13 +29,13 @@ func TestInitialSessionInputOfficialClient(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(worker)) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, worker)) if err != nil { t.Fatal(err) } server := httptest.NewServer(handler) defer server.Close() - unsupported, err := publicHandler(t, s, db, auth, "fake_alpha", workerExecution(worker)) + unsupported, err := publicHandler(t, s, db, auth, "fake_alpha", workerExecution(t, worker)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_metadata.go b/services/core/internal/store/session_metadata.go deleted file mode 100644 index 24ae3c2ff..000000000 --- a/services/core/internal/store/session_metadata.go +++ /dev/null @@ -1,46 +0,0 @@ -package store - -import ( - "context" - "errors" - "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/metadata" - "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/sessions" - "github.com/google/uuid" - "github.com/jackc/pgx/v5" -) - -func (s *Store) UpdateSessionMetadata(ctx context.Context, tenantID, sessionID string, values map[string]string) (sessions.Session, error) { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.Session{}, err - } - id := pgunit.PathID(sessionID) - encoded, err := metadata.Encode(values) - if err != nil { - return sessions.Session{}, fmt.Errorf("%w: %w", sessions.ErrInvalidInput, err) - } - var row sqlc.Session - err = s.pooled.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - var err error - row, err = q.UpdateSessionMetadata(ctx, sqlc.UpdateSessionMetadataParams{TenantID: tenant, ID: id, Metadata: encoded}) - if err != nil { - return err - } - return auditpg.RecordWriteAudit(ctx, q, tenantID, "update", "session", uuid.UUID(row.ID.Bytes).String(), "") - }) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.Session{}, sessions.ErrNotFound - } - if err != nil { - return sessions.Session{}, fmt.Errorf("update session metadata: %w", err) - } - session, decodeErr := sessionpg.SessionFromRow(row) - return s.sessionActivity(ctx, session, decodeErr) -} diff --git a/services/core/internal/store/session_metadata_test.go b/services/core/internal/store/session_metadata_test.go index 4e775bda4..a506b7b2a 100644 --- a/services/core/internal/store/session_metadata_test.go +++ b/services/core/internal/store/session_metadata_test.go @@ -32,7 +32,7 @@ func TestSessionMetadataPreservesCreationAndExecutionData(t *testing.T) { } before := snapshot() for _, metadata := range []map[string]string{{"new": "value"}, nil, {}, {"unicode": "中文🧪"}} { - updated, err := s.UpdateSessionMetadata(ctx, tenant, first.ID, metadata) + updated, err := sessionService(t, s).UpdateSessionMetadata(ctx, sessions.UpdateSessionMetadataCommand{TenantID: tenant, SessionID: first.ID, Metadata: metadata}) if metadata == nil { metadata = map[string]string{} } @@ -52,7 +52,7 @@ func TestSessionMetadataPreservesCreationAndExecutionData(t *testing.T) { t.Fatalf("changed creation request: %v", err) } } - current, err := s.GetSession(ctx, tenant, first.ID) + current, err := sessionAdapter(s).GetSession(ctx, tenant, first.ID) if err != nil { t.Fatal(err) } @@ -68,11 +68,11 @@ func TestSessionMetadataPreservesCreationAndExecutionData(t *testing.T) { {tenant, "invalid", nil, sessions.ErrNotFound}, {tenant, first.ID, map[string]string{"large": strings.Repeat("x", 64*1024)}, sessions.ErrInvalidInput}, } { - if _, err := s.UpdateSessionMetadata(ctx, test.tenant, test.session, test.metadata); !errors.Is(err, test.want) { + if _, err := sessionService(t, s).UpdateSessionMetadata(ctx, sessions.UpdateSessionMetadataCommand{TenantID: test.tenant, SessionID: test.session, Metadata: test.metadata}); !errors.Is(err, test.want) { t.Fatalf("rejected update error = %v, want %v", err, test.want) } } - got, err := s.GetSession(ctx, tenant, first.ID) + got, err := sessionAdapter(s).GetSession(ctx, tenant, first.ID) if err != nil || !reflect.DeepEqual(got, current) { t.Fatalf("rejected update changed Session: %+v, %v", got, err) } @@ -92,14 +92,14 @@ func TestSessionMetadataConcurrentReplacement(t *testing.T) { go func() { defer wg.Done() key := fmt.Sprint(i) - got, err := s.UpdateSessionMetadata(ctx, tenant, first.ID, map[string]string{key: key}) + got, err := sessionService(t, s).UpdateSessionMetadata(ctx, sessions.UpdateSessionMetadataCommand{TenantID: tenant, SessionID: first.ID, Metadata: map[string]string{key: key}}) if err != nil || len(got.Metadata) != 1 || got.Metadata[key] != key { t.Errorf("concurrent update = %+v, %v", got, err) } }() } wg.Wait() - got, err := s.GetSession(ctx, tenant, first.ID) + got, err := sessionAdapter(s).GetSession(ctx, tenant, first.ID) if err != nil || len(got.Metadata) != 1 { t.Fatalf("concurrent replacements merged or lost metadata: %+v, %v", got, err) } @@ -124,12 +124,12 @@ func TestSessionMetadataPreservesTerminalActivity(t *testing.T) { if _, err := transitionTurn(ctx, s, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: status}); err != nil { t.Fatal(err) } - before, err := s.GetSession(ctx, tenant, session.ID) + before, err := sessionAdapter(s).GetSession(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } before.Metadata = map[string]string{"label": status} - updated, err := s.UpdateSessionMetadata(ctx, tenant, session.ID, before.Metadata) + updated, err := sessionService(t, s).UpdateSessionMetadata(ctx, sessions.UpdateSessionMetadataCommand{TenantID: tenant, SessionID: session.ID, Metadata: before.Metadata}) if err != nil || !reflect.DeepEqual(updated, before) { t.Fatalf("metadata changed %s activity: %+v, %v", status, updated, err) } diff --git a/services/core/internal/store/session_reference_retry_public_test.go b/services/core/internal/store/session_reference_retry_public_test.go index c2068948f..0b07d6852 100644 --- a/services/core/internal/store/session_reference_retry_public_test.go +++ b/services/core/internal/store/session_reference_retry_public_test.go @@ -32,7 +32,7 @@ func TestSavedReferenceRetryOfficialClient(t *testing.T) { t.Error(err) } }) - handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(worker)) + handler, err := publicHandler(t, s, db, auth, "codex", workerExecution(t, worker)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_transaction.go b/services/core/internal/store/session_transaction.go index ef0749ffb..dbf6fe7bd 100644 --- a/services/core/internal/store/session_transaction.go +++ b/services/core/internal/store/session_transaction.go @@ -22,19 +22,6 @@ func (s *Store) withPublicSession(ctx context.Context, tenantID, sessionID strin } func (s *Store) withSessionState(ctx context.Context, tenantID, sessionID string, public bool, apply func(context.Context, *sqlc.Queries, pgtype.UUID) error) error { - return s.withLockedSession(ctx, tenantID, sessionID, public, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, locked sessions.LockedSession) error { - if public { - if err := locked.Public(); err != nil { - return err - } - } - return apply(ctx, q, session) - }) -} - -// withLockedSession resolves the identifiers and runs apply in the Session -// transaction, including for a publicly deleted Session. -func (s *Store) withLockedSession(ctx context.Context, tenantID, sessionID string, public bool, apply func(context.Context, *sqlc.Queries, pgtype.UUID, sessions.LockedSession) error) error { tenant, err := parseID(tenantID) if err != nil { return err @@ -47,6 +34,11 @@ func (s *Store) withLockedSession(ctx context.Context, tenantID, sessionID strin } } return sessionpg.WithSession(ctx, s.writer, tenant, id, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { - return apply(ctx, q, id, locked) + if public { + if err := locked.Public(); err != nil { + return err + } + } + return apply(ctx, q, id) }) } diff --git a/services/core/internal/store/session_write_audit.go b/services/core/internal/store/session_write_audit.go deleted file mode 100644 index 4ab76b688..000000000 --- a/services/core/internal/store/session_write_audit.go +++ /dev/null @@ -1,22 +0,0 @@ -package store - -import ( - "context" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" - "github.com/jackc/pgx/v5/pgtype" -) - -// AuditSessionOperation records an authorized public no-op or creation replay. -// It cannot create ownership or admit execution work. -func (s *Store) AuditSessionOperation(ctx context.Context, tenantID, sessionID, action string) error { - if action != "create" && action != "send_events" { - return sessions.ErrInvalidInput - } - return s.withPublicSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - return auditpg.RecordWriteAudit(ctx, q, tenantID, action, "session", uuid.UUID(session.Bytes).String(), "") - }) -} diff --git a/services/core/internal/store/session_write_audit_test.go b/services/core/internal/store/session_write_audit_test.go index c72082ed9..fef29312b 100644 --- a/services/core/internal/store/session_write_audit_test.go +++ b/services/core/internal/store/session_write_audit_test.go @@ -60,7 +60,7 @@ func TestSessionWriteAuditCreationReplayNoopAndDeletion(t *testing.T) { if _, err := s.CreateSession(sessionAuditContext(t, tenant, "b"), tenant, input); err != nil { t.Fatal(err) } - if _, err := s.GetSession(sessionAuditContext(t, tenant, "b"), tenant, created.ID); err != nil { + if _, err := sessionAdapter(s).GetSession(sessionAuditContext(t, tenant, "b"), tenant, created.ID); err != nil { t.Fatal(err) } sessionAuditCount(t, s, tenant, 2) @@ -71,15 +71,15 @@ func TestSessionWriteAuditCreationReplayNoopAndDeletion(t *testing.T) { } } for _, action := range []string{"create", "send_events"} { - if err := s.AuditSessionOperation(sessionAuditContext(t, tenant, "b"), tenant, created.ID, action); err != nil { + if err := sessionService(t, s).AuditSessionOperation(sessionAuditContext(t, tenant, "b"), sessions.AuditSessionOperationCommand{TenantID: tenant, SessionID: created.ID, Action: action}); err != nil { t.Fatal(err) } } - if _, err := s.UpdateSessionMetadata(sessionAuditContext(t, tenant, "b"), tenant, created.ID, map[string]string{"private": "not in audit"}); err != nil { + if _, err := sessionService(t, s).UpdateSessionMetadata(sessionAuditContext(t, tenant, "b"), sessions.UpdateSessionMetadataCommand{TenantID: tenant, SessionID: created.ID, Metadata: map[string]string{"private": "not in audit"}}); err != nil { t.Fatal(err) } for range 2 { - if err := s.DeleteSession(sessionAuditContext(t, tenant, "b"), tenant, created.ID); err != nil { + if err := sessionService(t, s).DeleteSession(sessionAuditContext(t, tenant, "b"), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: created.ID}); err != nil { t.Fatal(err) } } @@ -88,7 +88,7 @@ func TestSessionWriteAuditCreationReplayNoopAndDeletion(t *testing.T) { if err := s.pool.QueryRow(t.Context(), "SELECT jsonb_agg(to_jsonb(o))::text FROM write_audit_operations o WHERE tenant_id=$1", tenant).Scan(&history); err != nil || strings.Contains(history, "not in audit") || strings.Contains(history, "/workspace") { t.Fatal("payload entered audit", err) } - if err := s.AuditSessionOperation(sessionAuditContext(t, tenant, "b"), tenant, created.ID, "send_events"); !errors.Is(err, sessions.ErrNotFound) { + if err := sessionService(t, s).AuditSessionOperation(sessionAuditContext(t, tenant, "b"), sessions.AuditSessionOperationCommand{TenantID: tenant, SessionID: created.ID, Action: "send_events"}); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted no-op accepted", err) } sessionAuditCount(t, s, tenant, 7) @@ -139,20 +139,20 @@ func TestSessionWriteAuditRollback(t *testing.T) { } switch operation { case "update": - _, err = s.UpdateSessionMetadata(ctx, tenant, created.ID, map[string]string{"new": "value"}) + _, err = sessionService(t, s).UpdateSessionMetadata(ctx, sessions.UpdateSessionMetadataCommand{TenantID: tenant, SessionID: created.ID, Metadata: map[string]string{"new": "value"}}) case "delete": - err = s.DeleteSession(ctx, tenant, created.ID) + err = sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: created.ID}) case "events": _, err = s.SubmitInputs(ctx, tenant, created.ID, "events", []sessions.Input{messageInput("private")}) case "noop": - err = s.AuditSessionOperation(ctx, tenant, created.ID, "send_events") + err = sessionService(t, s).AuditSessionOperation(ctx, sessions.AuditSessionOperationCommand{TenantID: tenant, SessionID: created.ID, Action: "send_events"}) case "reserve": _, err = s.ReserveEnvironmentInput(ctx, tenant, created.ID, "reserve", []sessions.Input{messageInput("private")}) } if err == nil { t.Fatal("mutation bypassed audit failure") } - got, err := s.GetSession(t.Context(), tenant, created.ID) + got, err := sessionAdapter(s).GetSession(t.Context(), tenant, created.ID) if err != nil || len(got.Metadata) != 0 { t.Fatal("resource update/deletion survived rollback", got, err) } @@ -215,7 +215,7 @@ func TestSessionWriteAuditInitialInputAndHistoricalReplay(t *testing.T) { if err != nil { t.Fatal(err) } - if err := s.AuditSessionOperation(sessionAuditContext(t, tenant, "b"), tenant, legacy.ID, "create"); err != nil { + if err := sessionService(t, s).AuditSessionOperation(sessionAuditContext(t, tenant, "b"), sessions.AuditSessionOperationCommand{TenantID: tenant, SessionID: legacy.ID, Action: "create"}); err != nil { t.Fatal(err) } var owners int diff --git a/services/core/internal/store/sessions.go b/services/core/internal/store/sessions.go index 28b66381a..95c89ba13 100644 --- a/services/core/internal/store/sessions.go +++ b/services/core/internal/store/sessions.go @@ -167,64 +167,18 @@ func (s *Store) createSession(ctx context.Context, tenantID string, input sessio return sessions.Creation{Session: session, Created: row.ID == params.ID, Cursor: row.EventSequence}, err } -// GetSession always scopes lookup to the authenticated caller's tenant. -func (s *Store) GetSession(ctx context.Context, tenantID, sessionID string) (sessions.Session, error) { - tenant, err := parseID(tenantID) +// sessionActivity adds the activity projection to a Session read outside a +// snapshot, in a snapshot of its own. +func (s *Store) sessionActivity(ctx context.Context, session sessions.Session, err error) (sessions.Session, error) { if err != nil { return sessions.Session{}, err } - id := pgunit.PathID(sessionID) - row, err := s.queries.GetSession(ctx, sqlc.GetSessionParams{TenantID: tenant, ID: id}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.Session{}, sessions.ErrNotFound - } - if err != nil { - return sessions.Session{}, fmt.Errorf("get session: %w", err) - } - session, decodeErr := sessionpg.SessionFromRow(row) - return s.sessionActivity(ctx, session, decodeErr) -} - -// ListSessions orders by creation time and ID. The cursor is the last returned -// session ID and must belong to the same tenant; it grants no additional access. -func (s *Store) ListSessions(ctx context.Context, tenantID, cursor string, limit int, ascending bool, agentID *string) (sessions.Page, error) { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.Page{}, err - } - if limit < 1 || limit > 100 { - return sessions.Page{}, fmt.Errorf("%w: page size must be 1..100", sessions.ErrInvalidInput) - } - params := sqlc.ListSessionsParams{TenantID: tenant, PageLimit: int32(limit + 1), AfterID: pgtype.UUID{Valid: true}, Ascending: ascending} - if agentID != nil { - params.AgentID = pgtype.Text{String: *agentID, Valid: true} - } - if cursor != "" { - after, err := s.GetSession(ctx, tenantID, pgunit.LookupCursor(cursor)) - if err != nil { - return sessions.Page{}, err - } - params.AfterCreated = pgtype.Timestamptz{Time: after.CreatedAt, Valid: true} - params.AfterID, _ = parseID(after.ID) - } - rows, err := s.queries.ListSessions(ctx, params) - if err != nil { - return sessions.Page{}, fmt.Errorf("list sessions: %w", err) - } - page := sessions.Page{Sessions: make([]sessions.Session, 0, min(limit, len(rows)))} - if len(rows) > limit { - page.NextCursor = uuid.UUID(rows[limit-1].ID.Bytes).String() - rows = rows[:limit] - } - for _, row := range rows { - session, err := sessionpg.SessionFromRow(row) - session, err = s.sessionActivity(ctx, session, err) - if err != nil { - return sessions.Page{}, err - } - page.Sessions = append(page.Sessions, session) - } - return page, nil + err = s.pooled.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { + var err error + session, err = sessionpg.LoadSessionActivity(ctx, s.queries.WithTx(tx), session) + return err + }) + return session, err } // parseID translates pgunit's identifier rule into store's invalid-input diff --git a/services/core/internal/store/sessions_test.go b/services/core/internal/store/sessions_test.go index 09d4bdd62..076938e3f 100644 --- a/services/core/internal/store/sessions_test.go +++ b/services/core/internal/store/sessions_test.go @@ -65,10 +65,10 @@ func TestSessionsPersistAndStayTenantScoped(t *testing.T) { if first.ID == other.ID { t.Fatal("idempotency leaked across tenants") } - if _, err := s.GetSession(ctx, tenantB, first.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).GetSession(ctx, tenantB, first.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("cross-tenant read: %v", err) } - if _, err := s.ListSessions(ctx, tenantB, first.ID, 10, false, nil); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionAdapter(s).ListSessions(ctx, tenantB, first.ID, 10, false, nil); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("cross-tenant cursor: %v", err) } for _, key := range []string{"second", "third"} { @@ -80,14 +80,14 @@ func TestSessionsPersistAndStayTenantScoped(t *testing.T) { // Recreate the pool and Store as a new service process would. pool.Close() recovered, _ := testStore(t) - got, err := recovered.GetSession(ctx, tenantA, first.ID) + got, err := sessionAdapter(recovered).GetSession(ctx, tenantA, first.ID) if err != nil || !reflect.DeepEqual(got, first) { t.Fatalf("restart read = %+v, %v; want %+v", got, err, first) } seen := map[string]bool{} cursor := "" for { - page, err := recovered.ListSessions(ctx, tenantA, cursor, 2, false, nil) + page, err := sessionAdapter(recovered).ListSessions(ctx, tenantA, cursor, 2, false, nil) if err != nil { t.Fatal(err) } @@ -108,7 +108,7 @@ func TestSessionsPersistAndStayTenantScoped(t *testing.T) { if len(seen) != 3 || !seen[first.ID] || seen[other.ID] { t.Fatalf("pagination lost or leaked sessions: %+v", seen) } - empty, err := recovered.ListSessions(ctx, uuid.NewString(), "", 10, false, nil) + empty, err := sessionAdapter(recovered).ListSessions(ctx, uuid.NewString(), "", 10, false, nil) if err != nil || empty.Sessions == nil || len(empty.Sessions) != 0 { t.Fatalf("empty tenant = %+v, %v", empty, err) } @@ -159,7 +159,7 @@ func TestConcurrentSessionCreationIsIdempotent(t *testing.T) { t.Fatalf("changed request = %v", err) } } - page, err := s.ListSessions(ctx, tenant, "", 10, false, nil) + page, err := sessionAdapter(s).ListSessions(ctx, tenant, "", 10, false, nil) if err != nil || len(page.Sessions) != 1 || !reflect.DeepEqual(page.Sessions[0], replay) { t.Fatalf("retry changed stored session: %+v, %v", page, err) } diff --git a/services/core/internal/store/structured_output_native_test.go b/services/core/internal/store/structured_output_native_test.go index 2d51e2995..e92633af3 100644 --- a/services/core/internal/store/structured_output_native_test.go +++ b/services/core/internal/store/structured_output_native_test.go @@ -45,7 +45,7 @@ func TestNativeStructuredOutputPublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, "claude_sdk", workerExecution(worker), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "claude_sdk", workerExecution(t, worker), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/subagent_identities_test.go b/services/core/internal/store/subagent_identities_test.go index f5dc4a1a5..07aa5228f 100644 --- a/services/core/internal/store/subagent_identities_test.go +++ b/services/core/internal/store/subagent_identities_test.go @@ -58,7 +58,7 @@ func TestSubagentIdentityIsAtomicScopedAndImmutable(t *testing.T) { t.Fatal("foreign write", err) } } - before, err := s.SessionEventCursor(ctx, tenant, session.ID) + before, err := sessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } @@ -80,7 +80,7 @@ func TestSubagentIdentityIsAtomicScopedAndImmutable(t *testing.T) { if err != nil || len(events) != 3 { t.Fatal("partial journal survived", len(events), err) } - cursor, err := s.SessionEventCursor(ctx, tenant, session.ID) + cursor, err := sessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) if err != nil || cursor != before { t.Fatal("partial public projection survived", cursor, err) } @@ -132,7 +132,7 @@ func TestSubagentIdentityIsAtomicScopedAndImmutable(t *testing.T) { if err != nil || !reflect.DeepEqual(again, saved) { t.Fatal("continuation changed immutable first observation", again, err) } - if err = reopened.DeleteSession(ctx, tenant, session.ID); !errors.Is(err, sessions.ErrNotIdle) { + if err = sessionService(t, reopened).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("running Session deleted", err) } if err = reopened.commitLegacyDeletion(ctx, tenant, session.ID); err != nil { diff --git a/services/core/internal/store/subagent_native_outputs_test.go b/services/core/internal/store/subagent_native_outputs_test.go index f6154fb3b..94511c8e9 100644 --- a/services/core/internal/store/subagent_native_outputs_test.go +++ b/services/core/internal/store/subagent_native_outputs_test.go @@ -51,7 +51,7 @@ func TestSubagentNativeFunctionResultDoesNotConsumeOutputIndex(t *testing.T) { t.Fatal(items, err) } // Child Items publish no Session events; only root work reaches the stream. - events, err := s.ListSessionEvents(t.Context(), tenant, session.ID, 0) + events, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, 0) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/subagent_resources_test.go b/services/core/internal/store/subagent_resources_test.go index 87135d8a1..fed5afe72 100644 --- a/services/core/internal/store/subagent_resources_test.go +++ b/services/core/internal/store/subagent_resources_test.go @@ -151,7 +151,7 @@ func TestSubagentResourcesNativeOwnershipLifecycleAndRecovery(t *testing.T) { if err != nil || value.Status != "active" || value.ClosedAt != nil || value.OpenedAt != 100 { t.Fatal(value, err) } - changes, err := s.ListSessionEvents(ctx, tenant, session.ID, 0) + changes, err := sessionAdapter(s).ListSessionEvents(ctx, tenant, session.ID, 0) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/template_composition_public_test.go b/services/core/internal/store/template_composition_public_test.go index 1585ea31b..078d99152 100644 --- a/services/core/internal/store/template_composition_public_test.go +++ b/services/core/internal/store/template_composition_public_test.go @@ -77,7 +77,7 @@ func TestTemplateCompositionOfficialClientPostgres(t *testing.T) { } t.Cleanup(func() { for _, id := range receipt.Sessions { - if err := s.DeleteSession(context.Background(), tenant, id); err != nil { + if err := store.SessionService(t, s).DeleteSession(context.Background(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: id}); err != nil { t.Error(err) } } diff --git a/services/core/internal/store/template_null_selection_public_test.go b/services/core/internal/store/template_null_selection_public_test.go index edfe2a46e..849cd302f 100644 --- a/services/core/internal/store/template_null_selection_public_test.go +++ b/services/core/internal/store/template_null_selection_public_test.go @@ -82,7 +82,7 @@ func TestTemplateNullSelectionOfficialClientPostgres(t *testing.T) { } t.Cleanup(func() { for _, id := range receipt.Sessions { - if err := s.DeleteSession(context.Background(), tenant, id); err != nil { + if err := store.SessionService(t, s).DeleteSession(context.Background(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: id}); err != nil { t.Error(err) } } diff --git a/services/core/internal/store/token_usage.go b/services/core/internal/store/token_usage.go deleted file mode 100644 index e5a57312f..000000000 --- a/services/core/internal/store/token_usage.go +++ /dev/null @@ -1,30 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" -) - -// MeasuredSessionUsage returns Core-internal measured usage for Runtime -// telemetry: the sum of every recorded root Turn snapshot, active Turns -// included, and null only when nothing is recorded. Public Session usage keeps -// the official rule of SessionTokenUsage. A missing Session reads as null, so -// callers resolve the Session first. -func (s *Store) MeasuredSessionUsage(ctx context.Context, tenantID, sessionID string) (json.RawMessage, error) { - tenant, err := parseID(tenantID) - if err != nil { - return nil, err - } - id, err := parseID(sessionID) - if err != nil { - return nil, err - } - usage, err := s.queries.SessionMeasuredTokenUsage(ctx, sqlc.SessionMeasuredTokenUsageParams{TenantID: tenant, ID: id}) - if err != nil { - return nil, fmt.Errorf("read measured session usage: %w", err) - } - return usage, nil -} diff --git a/services/core/internal/store/token_usage_integration_test.go b/services/core/internal/store/token_usage_integration_test.go index 2dd7268ff..a042c21c2 100644 --- a/services/core/internal/store/token_usage_integration_test.go +++ b/services/core/internal/store/token_usage_integration_test.go @@ -85,7 +85,7 @@ func TestTokenUsageDurableSnapshotsAndSessionTotals(t *testing.T) { } defer restored.Close() fresh := store.New(restored) - got, err := fresh.GetSession(ctx, tenant, session.ID) + got, err := store.SessionAdapter(fresh).GetSession(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } @@ -96,11 +96,11 @@ func TestTokenUsageDurableSnapshotsAndSessionTotals(t *testing.T) { if total.InputTokens != 40 || total.OutputTokens != 6 || total.TotalTokens != 46 || total.InputTokensDetails.CachedTokens != 8 || total.OutputTokensDetails.ReasoningTokens != 4 { t.Fatalf("double counted totals: %+v", total) } - page, err := fresh.ListSessions(ctx, tenant, "", 100, true, nil) + page, err := store.SessionAdapter(fresh).ListSessions(ctx, tenant, "", 100, true, nil) if err != nil || len(page.Sessions) != 1 || string(page.Sessions[0].Usage) != string(got.Usage) { t.Fatalf("list totals: %+v %v", page, err) } - if _, err = fresh.GetSession(ctx, uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err = store.SessionAdapter(fresh).GetSession(ctx, uuid.NewString(), session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } } @@ -156,22 +156,22 @@ func TestSessionUsageRequiresEveryRootTurnEndedAndMeasured(t *testing.T) { // included, and is scoped to the tenant. measured := func(want int64) { t.Helper() - got, err := s.MeasuredSessionUsage(ctx, tenant, session.ID) + got, err := store.SessionAdapter(s).MeasuredSessionUsage(ctx, tenant, session.ID) var value v1.TokenUsage if err != nil || (want < 0) != (got == nil) || (want >= 0 && (json.Unmarshal(got, &value) != nil || value.TotalTokens != want)) { t.Fatalf("measured usage = %s %v, want total %d", got, err, want) } - if foreign, err := s.MeasuredSessionUsage(ctx, uuid.NewString(), session.ID); err != nil || foreign != nil { + if foreign, err := store.SessionAdapter(s).MeasuredSessionUsage(ctx, uuid.NewString(), session.ID); err != nil || foreign != nil { t.Fatalf("foreign measured usage: %s %v", foreign, err) } } total := func(want int64) { t.Helper() - got, err := s.GetSession(ctx, tenant, session.ID) + got, err := store.SessionAdapter(s).GetSession(ctx, tenant, session.ID) if err != nil { t.Fatal(err) } - page, err := s.ListSessions(ctx, tenant, "", 100, true, nil) + page, err := store.SessionAdapter(s).ListSessions(ctx, tenant, "", 100, true, nil) if err != nil || len(page.Sessions) != 1 || string(page.Sessions[0].Usage) != string(got.Usage) { t.Fatalf("list usage: %+v %v", page, err) } @@ -208,7 +208,7 @@ func TestSessionUsageRequiresEveryRootTurnEndedAndMeasured(t *testing.T) { } lastIdleUsage := func() json.RawMessage { t.Helper() - changes, err := s.ListSessionEvents(ctx, tenant, session.ID, 0) + changes, err := store.SessionAdapter(s).ListSessionEvents(ctx, tenant, session.ID, 0) if err != nil || len(changes) == 0 || changes[len(changes)-1].Event.Type != "agent.session.idle" { t.Fatal(changes, err) } diff --git a/services/core/internal/store/tool_policy_native_test.go b/services/core/internal/store/tool_policy_native_test.go index 821c075cf..bfa83bfc0 100644 --- a/services/core/internal/store/tool_policy_native_test.go +++ b/services/core/internal/store/tool_policy_native_test.go @@ -51,7 +51,7 @@ func TestNativeToolPolicyPublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: foreignTenant, SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: foreignTenant}, }) - handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, kind, workerExecution(t, worker), withPolicy(h.d.Policy), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } @@ -77,17 +77,17 @@ func TestNativeToolPolicyPublicExecution(t *testing.T) { if err != nil || json.Unmarshal(raw, &proof) != nil || len(proof.Sessions) != 4 { t.Fatal("invalid public evidence", err) } - page, err := h.s.ListSessions(ctx, h.tenant, "", 100, true, nil) + page, err := store.SessionAdapter(h.s).ListSessions(ctx, h.tenant, "", 100, true, nil) if err != nil || len(page.Sessions) != 1+len(proof.Sessions) { t.Fatal("rejected configuration persisted a Session", err) } - foreignPage, err := h.s.ListSessions(ctx, foreignTenant, "", 100, true, nil) + foreignPage, err := store.SessionAdapter(h.s).ListSessions(ctx, foreignTenant, "", 100, true, nil) if err != nil || len(foreignPage.Sessions) != 0 { t.Fatal("foreign Agent reference persisted a Session", err) } nativeIDs := make(map[string]string, len(proof.Sessions)) for _, item := range proof.Sessions { - session, err := h.s.GetSession(ctx, h.tenant, item.ID) + session, err := store.SessionAdapter(h.s).GetSession(ctx, h.tenant, item.ID) if err != nil || session.Engine != kind { t.Fatal("selected engine was not persisted", err) } diff --git a/services/core/internal/store/tool_search_native_test.go b/services/core/internal/store/tool_search_native_test.go index 626443f9b..b47764be6 100644 --- a/services/core/internal/store/tool_search_native_test.go +++ b/services/core/internal/store/tool_search_native_test.go @@ -45,7 +45,7 @@ func TestNativeToolSearchPublicExecution(t *testing.T) { {OrganizationID: "test", ProjectID: h.tenant, SubjectKind: "service_account", SubjectID: "owner", TokenSHA256: runtimedevice.HashCredential(token), TenantID: h.tenant}, {OrganizationID: "test", ProjectID: uuid.NewString(), SubjectKind: "service_account", SubjectID: "other", TokenSHA256: runtimedevice.HashCredential(foreign), TenantID: uuid.NewString()}, }) - handler, err := publicHandler(t, h.s, h.db, auth, "claude_sdk", workerExecution(worker), nativeDeploymentDefaults(model, provider)) + handler, err := publicHandler(t, h.s, h.db, auth, "claude_sdk", workerExecution(t, worker), nativeDeploymentDefaults(model, provider)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/turn_inputs_test.go b/services/core/internal/store/turn_inputs_test.go index c5924e554..3ffc28e26 100644 --- a/services/core/internal/store/turn_inputs_test.go +++ b/services/core/internal/store/turn_inputs_test.go @@ -160,7 +160,7 @@ func TestTurnInputRetriesAndRestart(t *testing.T) { if len(all) != 2 || all[0].Sequence != first.Sequence || all[1].Sequence != steer.Sequence { t.Fatalf("recovered inputs = %+v", all) } - snapshot, err := recovered.GetSession(ctx, tenant, session.ID) + snapshot, err := sessionAdapter(recovered).GetSession(ctx, tenant, session.ID) if err != nil || snapshot.LastTurn == nil || snapshot.LastTurn.ID != next.TurnID || snapshot.LastTurn.Status != sessions.TurnQueued { t.Fatal("latest Session activity did not survive restart", err) } diff --git a/services/core/internal/store/unified_model_configuration_http_test.go b/services/core/internal/store/unified_model_configuration_http_test.go index ad3b4ae6c..ef3f4e6c3 100644 --- a/services/core/internal/store/unified_model_configuration_http_test.go +++ b/services/core/internal/store/unified_model_configuration_http_test.go @@ -106,7 +106,7 @@ func TestUnifiedModelConfigurationHTTP(t *testing.T) { } } equalJSON(publicNative, native) - snapshot, err := st.GetSessionExecutionConfiguration(t.Context(), tenant, id) + snapshot, err := store.SessionAdapter(st).GetSessionExecutionConfiguration(t.Context(), tenant, id) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/whitespace_input_public_test.go b/services/core/internal/store/whitespace_input_public_test.go index d03252e9f..7ea243498 100644 --- a/services/core/internal/store/whitespace_input_public_test.go +++ b/services/core/internal/store/whitespace_input_public_test.go @@ -106,7 +106,7 @@ func TestWhitespaceOnlyTextHarnessAdmissionPostgres(t *testing.T) { } }) serve := func(engine string) pathIDClient { - handler, err := publicHandler(t, s, db, auth, engine, workerExecution(worker), executorURL("https://offline-executor.example")) + handler, err := publicHandler(t, s, db, auth, engine, workerExecution(t, worker), executorURL("https://offline-executor.example")) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/worker_preparation_failure_test.go b/services/core/internal/store/worker_preparation_failure_test.go index 7c9dfd50b..93b293053 100644 --- a/services/core/internal/store/worker_preparation_failure_test.go +++ b/services/core/internal/store/worker_preparation_failure_test.go @@ -7,6 +7,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestWorkerSettlesConfirmedPreparationFailureAndAcceptsNewInput(t *testing.T) { @@ -33,11 +34,11 @@ func TestWorkerSettlesConfirmedPreparationFailureAndAcceptsNewInput(t *testing.T current, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) return err == nil && current.State == sessions.EnvironmentInputFailed }) - session, err := h.s.GetSession(t.Context(), h.tenant, h.session.ID) + session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) if err != nil || session.PendingInput || session.LastTurn != nil || session.EnvironmentInputActivity == nil || session.EnvironmentInputActivity.Failure != "runtime_preparation_failed" { t.Fatal("preparation did not release input with a safe failure", err) } - changes, err := h.s.ListSessionEvents(t.Context(), h.tenant, h.session.ID, 0) + changes, err := store.SessionAdapter(h.s).ListSessionEvents(t.Context(), h.tenant, h.session.ID, 0) if err != nil { t.Fatal(err) } From e8cc52d5e92e5c153d5e18cd7348c27f6cd59d15 Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 02:48:39 +0800 Subject: [PATCH 5/6] Move the administrator summary into sessions (#439) --- contracts/agents-api/core.openapi.yaml | 32 ++++---- services/core/IMPLEMENTATION.md | 4 +- services/core/cmd/server/http_routes_test.go | 2 +- services/core/cmd/server/main.go | 2 +- services/core/internal/api/admin_resources.go | 6 +- .../core/internal/api/admin_resources_test.go | 9 +-- services/core/internal/api/admin_runtime.go | 9 +-- .../core/internal/api/admin_runtime_test.go | 2 +- services/core/internal/api/admin_summary.go | 33 ++++---- services/core/internal/api/dependencies.go | 3 +- .../core/internal/api/dependencies_test.go | 5 +- services/core/internal/api/fakes_test.go | 15 ++-- .../persistence/postgres/sessionpg/admin.go | 52 ++++++++++++ .../postgres/sessionpg/admin_test.go | 44 ++++++++++ .../providers/configuration_flow_test.go | 2 +- services/core/internal/sessions/admin.go | 24 +++++- services/core/internal/store/admin_summary.go | 81 ------------------- .../store/public_handler_fixture_test.go | 2 +- 18 files changed, 174 insertions(+), 153 deletions(-) delete mode 100644 services/core/internal/store/admin_summary.go diff --git a/contracts/agents-api/core.openapi.yaml b/contracts/agents-api/core.openapi.yaml index ddf6a108d..9a74d1381 100644 --- a/contracts/agents-api/core.openapi.yaml +++ b/contracts/agents-api/core.openapi.yaml @@ -196,7 +196,7 @@ definitions: agent_id: type: string assets: - $ref: '#/definitions/store.AdminAssetCounts' + $ref: '#/definitions/sessions.AdminAssetCounts' coverage: $ref: '#/definitions/api.AdminUsageCoverage' key_id: @@ -1395,6 +1395,21 @@ definitions: source_commit: type: string type: object + sessions.AdminAssetCounts: + properties: + agents: + type: integer + credentials: + type: integer + environment_templates: + type: integer + files: + type: integer + skills: + type: integer + vaults: + type: integer + type: object sessions.ExecutorCredential: properties: created_at: @@ -1424,21 +1439,6 @@ definitions: state: type: string type: object - store.AdminAssetCounts: - properties: - agents: - type: integer - credentials: - type: integer - environment_templates: - type: integer - files: - type: integer - skills: - type: integer - vaults: - type: integer - type: object v1.Agent: properties: id: diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 338a330bc..3c5e0278d 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, `LoadRequiredActions` lists the calls a Turn waits on, and `TransitionTurn` moves a root Turn's status as its current status allows and settles a Turn that ends. `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 and Turn operation loads the Turn 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` or `TransitionTurn`, 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; among them `LoadSessionActivity`, the one Session activity read, adds the activity projection to the Sessions that `store`'s Session creation and administrator summary return. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; `placementpg`, which loads the placement facts and reserves placements, sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. +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` or `TransitionTurn`, 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; among them `LoadSessionActivity`, the one Session activity read, adds the activity projection to the Session that `store`'s Session creation returns. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; `placementpg`, which loads the placement facts and reserves placements, sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. `store` is transitional. `store.New` builds a pooled Store, and `store.NewExecution(s, lease)` builds the execution writer on a lease it borrows. An execution-only operation on a pooled Store fails with `store.ErrExecutionAuthority`. New adapters do not copy that check: their execution repositories require a `*pgunit.Lease` at construction, their public repositories expose no execution operation, and the check goes away with `store`. @@ -36,7 +36,7 @@ Domain owners, each with its PostgreSQL adapter under `internal/persistence/post - `deployment` and its subpackage `deployment/placement` (`deploymentpg`, `placementpg`): the sandbox deployment and its nodes: provider configuration and the sealed credential, the specification and retained generations, setup, update and switch, node enrollment, identity and authentication, generation configuration, capacity, presence, status, host history, reset, and the counts of nodes and sandboxes bound to the public address; and hosted runtime allocations: reservation with the dedicated device, compute ownership and settlement, observation diagnostics, activity, compute phases, wake receipts and cleanup, with the reads that schedule, discover and authorize them. It interprets Sandbox Provider declarations through the `providers.Registry` it is given, whose lookups return typed errors. `cmd/server` builds that registry and calls it only to build a direct Provider and to discover a Provider's configuration; it takes each setup's mode and declared operations from the `Setup` that `deployment` returns. `deployment/placement` owns hosted admission and placement: `cmd/server` builds one `placement.Rules` from the registry and the public URL, both fixed while Core runs, and gives it to `deployment.Service` and to `store`; its pure decisions admit a hosted Session, choose its node, and admit an allocation's reserved node and a restore on it, and its errors keep one status, code and message through `writeStoreError` and `writeDeploymentError`. `placementpg` loads the facts those decisions read and applies them on the caller's transaction-bound queries; it has no Store or transaction runner. The only other reader is `store`: Session creation decides admission and placement with those rules and `placementpg` inside its own transaction. Deployment changes and allocation writes run through `deployment.ExecutionOperations` on `deploymentpg.NewExecution(lease, …)`, which the Worker receives as `execution.Owner.Deployment`. Each allocation write locks the owning Session, decides in `deployment`, applies in `deploymentpg` and prunes the Session's change journal in the same leased transaction; cleanup settles the Session through the `sessions` procedures on a `sessionpg.SessionTx` bound to that transaction. The Worker reads the deployment, prepares a selection's setup and records live-compute activity through the pooled `deployment.Service` in `execution.Dispatcher.Deployment`, and lists the Sessions a reset still has to archive, schedules node lifecycles and reads allocations through the pooled `deployment.Reader` in `execution.Dispatcher.DeploymentReader`; a pooled read carries no lease, and the leased write that follows it rechecks the allocation's owner. Node management and reads use the pooled `deploymentpg.Store`, which `cmd/server` also reads the owner epoch from and runs the host-history sampler on. Provider calls run outside transactions, and the final transaction rechecks the expected generation. Session archive crosses the Session and the deployment, so it is a `deployment.ExecutionOperations` operation: `ArchiveSession`, which the administrator's archive route calls through `api.Execution.SessionArchive`, and `ArchiveResetSession`, which a reset calls, lock the Session through `deploymentpg`'s `WithSessionArchive`, check the deployment's generation and the running reset in `deployment`, then expire the Environment and cancel its work through the `sessions` procedures, revoke the allocation's device and request its cleanup, and record the audit entry in one leased transaction. `deployment.ObservationResolver` resolves a Session's Runtime observation target for `runtimeobs` from `sessions.SessionReader` and `deployment.Reader`. - `coremetrics` (`coremetricspg`): the Core metrics PostgreSQL holds: the root Turn queue counts, the root Turn history, read from one read-only snapshot, and the database size. `cmd/server`'s Core metrics source adds them to the process, pool, Worker and daemon registry measurements. - `runtimehistory` (`runtimehistory/postgresreader`, outside this directory): Runtime history samples: the periodic export, scoped reads and retention, which also prunes node-host samples. `runtimehistory.Service` scopes a read to the Session's hosted Environment through `sessions.EnvironmentReader`. -- `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 Sessions (their reads, the change journal and stream snapshot, diagnostics, the frozen execution configuration, measured usage and archive state, metadata updates, deletion and the public write audit), Turns (their reads and the execution work scan), the model provider a Session froze, which `sessionpg.Store` opens with the credential key, 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, which include a Session's execution binding and the execution device list, their creation, revocation, heartbeats and Runtime enrollment, and the archived cancellation receipt the daemon gateway reads), the Sessions whose Runtime the administrator observes across Projects, 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. Turn transitions, execution completion, which alone publishes or discards a Turn's staged Artifacts, the start of Artifact capture, 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`. +- `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 Sessions (their reads, the change journal and stream snapshot, diagnostics, the frozen execution configuration, measured usage and archive state, metadata updates, deletion and the public write audit), Turns (their reads and the execution work scan), the model provider a Session froze, which `sessionpg.Store` opens with the credential key, 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, which include a Session's execution binding and the execution device list, their creation, revocation, heartbeats and Runtime enrollment, and the archived cancellation receipt the daemon gateway reads), the administrator's views across Projects (a Project's asset counts and Sessions for the summary, and the Sessions whose Runtime the administrator observes), 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. Turn transitions, execution completion, which alone publishes or discards a Turn's staged Artifacts, the start of Artifact capture, function calls and their application receipts, the Turn execution journal, Environment initialization, connection observations and their reconciliation, Session device binding, and file-write reservation and settlement run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. ## Request handling diff --git a/services/core/cmd/server/http_routes_test.go b/services/core/cmd/server/http_routes_test.go index 05585fd38..f31600030 100644 --- a/services/core/cmd/server/http_routes_test.go +++ b/services/core/cmd/server/http_routes_test.go @@ -106,7 +106,7 @@ func daemonComposition(t testing.TB) http.Handler { Artifacts: struct{ api.Artifacts }{}, ArtifactsReader: struct{ api.ArtifactsReader }{}, SessionAdmin: struct{ api.SessionAdmin }{}, Environments: struct{ api.Environments }{}, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, ExecutorConnections: struct{ api.ExecutorConnections }{}, - Admin: struct{ api.Admin }{}, AdminRuntimeTargets: struct{ api.AdminRuntimeTargets }{}, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, Metrics: struct{ api.Metrics }{}, + Admin: struct{ api.Admin }{}, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, Metrics: struct{ api.Metrics }{}, RuntimeObservations: struct{ api.RuntimeObservations }{}, RuntimeHistory: struct{ api.RuntimeHistory }{}, Execution: &api.Execution{ ExecutorURL: "wss://core.example/api/v1/agent-daemon/ws", diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index 3cf6979dc..4eef9799e 100644 --- a/services/core/cmd/server/main.go +++ b/services/core/cmd/server/main.go @@ -409,7 +409,7 @@ func run() error { ArtifactsReader: sessionStore, SessionAdmin: sessionStore, Environments: sessionService, EnvironmentsReader: sessionStore, ExecutorConnections: executorConnections{sessions: sessionStore, registry: registry}, - Admin: executionStore, AdminRuntimeTargets: sessionStore, AdminAudit: auditStore, WriteAudit: auditStore, Metrics: metrics, + Admin: sessionStore, AdminAudit: auditStore, WriteAudit: auditStore, Metrics: metrics, RuntimeObservations: observationService, RuntimeHistory: historyService, } if worker != nil { diff --git a/services/core/internal/api/admin_resources.go b/services/core/internal/api/admin_resources.go index 79ad09fe7..c61ad30c3 100644 --- a/services/core/internal/api/admin_resources.go +++ b/services/core/internal/api/admin_resources.go @@ -5,7 +5,6 @@ import ( "net/http" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/go-chi/chi/v5" ) @@ -13,9 +12,10 @@ import ( // Only Core-key-authenticated project resource handlers receive it. type adminTenantContextKey struct{} -// Admin reads the administrator's asset summary of a Project. +// Admin reads the administrator's cross-Project Session views. type Admin interface { - ReadAdminSummary(context.Context, string, store.AdminSummaryFilter, func(sessions.Session, *string) error) (store.AdminAssetCounts, error) + ReadAdminSummary(context.Context, string, sessions.AdminSummaryFilter, func(sessions.Session, *string) error) (sessions.AdminAssetCounts, error) + ListAdminRuntimeTargets(context.Context, []string, string, int, bool) (sessions.AdminRuntimeTargetPage, error) } func (h *Handler) adminResourceScope(next http.Handler) http.Handler { diff --git a/services/core/internal/api/admin_resources_test.go b/services/core/internal/api/admin_resources_test.go index 2a64b778a..cf323c93c 100644 --- a/services/core/internal/api/admin_resources_test.go +++ b/services/core/internal/api/admin_resources_test.go @@ -13,7 +13,6 @@ import ( "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/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) const managementProjectID = "22222222-2222-4222-8222-222222222222" @@ -117,10 +116,10 @@ func TestAdminResourcesHaveExplicitTargetWithoutCallerImpersonation(t *testing.T type summaryFixture struct { tenant string - filter store.AdminSummaryFilter + filter sessions.AdminSummaryFilter } -func (s *summaryFixture) ReadAdminSummary(_ context.Context, tenant string, filter store.AdminSummaryFilter, visit func(sessions.Session, *string) error) (store.AdminAssetCounts, error) { +func (s *summaryFixture) ReadAdminSummary(_ context.Context, tenant string, filter sessions.AdminSummaryFilter, visit func(sessions.Session, *string) error) (sessions.AdminAssetCounts, error) { s.tenant, s.filter = tenant, filter for i, usage := range []json.RawMessage{nil, json.RawMessage(`{"input_tokens":3,"output_tokens":5,"total_tokens":8,"input_tokens_details":{"cached_tokens":2},"output_tokens_details":{"reasoning_tokens":1}}`)} { session := sessions.Session{ID: "session", TenantID: tenant, Configuration: json.RawMessage(`{"agent":{"id":"agent","model":"model","tools":[]},"environment":{"type":"none"}}`), CreatedAt: time.Unix(100+int64(i), 0), Usage: usage} @@ -128,10 +127,10 @@ func (s *summaryFixture) ReadAdminSummary(_ context.Context, tenant string, filt session.LastTurn = &sessions.Turn{Status: sessions.TurnInProgress, CreatedAt: time.Unix(110, 0)} } if err := visit(session, nil); err != nil { - return store.AdminAssetCounts{}, err + return sessions.AdminAssetCounts{}, err } } - return store.AdminAssetCounts{Agents: 4, Skills: 2}, nil + return sessions.AdminAssetCounts{Agents: 4, Skills: 2}, nil } func TestAdminSummaryUsesPublicStateAndNullUsageCoverage(t *testing.T) { key := callerBinding() diff --git a/services/core/internal/api/admin_runtime.go b/services/core/internal/api/admin_runtime.go index c0ff4716e..31d204105 100644 --- a/services/core/internal/api/admin_runtime.go +++ b/services/core/internal/api/admin_runtime.go @@ -7,15 +7,8 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimeobs" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -// AdminRuntimeTargets lists the Sessions whose Runtime the administrator -// observes across managed Projects. -type AdminRuntimeTargets interface { - ListAdminRuntimeTargets(context.Context, []string, string, int, bool) (sessions.AdminRuntimeTargetPage, error) -} - type AdminRuntimeObservation struct { ProjectID string `json:"project_id"` Observation AdminRuntimeObservationDetail `json:"observation"` @@ -79,7 +72,7 @@ func (h *Handler) adminRuntimeObservations(w http.ResponseWriter, r *http.Reques break } } - page, err := h.AdminRuntimeTargets.ListAdminRuntimeTargets(ctx, tenants, options.after, options.limit, options.ascending) + page, err := h.Admin.ListAdminRuntimeTargets(ctx, tenants, options.after, options.limit, options.ascending) if err != nil { writeSessionsError(w, r, err) return diff --git a/services/core/internal/api/admin_runtime_test.go b/services/core/internal/api/admin_runtime_test.go index 9df4f2579..d2962247a 100644 --- a/services/core/internal/api/admin_runtime_test.go +++ b/services/core/internal/api/admin_runtime_test.go @@ -39,7 +39,7 @@ func adminRuntimeFixture(t *testing.T, catalog []projects.Project, targets []ses return projects.Page{Data: catalog}, nil } management := &adminRuntimeTargets{page: sessions.AdminRuntimeTargetPage{Data: targets, HasMore: true}} - fakes.adminRuntimeTargets.listAdminRuntimeTargets = management.ListAdminRuntimeTargets + fakes.admin.listAdminRuntimeTargets = management.ListAdminRuntimeTargets observeWith(service)(&deps, fakes) return newTestHandler(t, deps), management } diff --git a/services/core/internal/api/admin_summary.go b/services/core/internal/api/admin_summary.go index 9805e7517..d24efa7c0 100644 --- a/services/core/internal/api/admin_summary.go +++ b/services/core/internal/api/admin_summary.go @@ -9,7 +9,6 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/projects" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) type AdminSessionCounts struct { @@ -25,14 +24,14 @@ type AdminUsageCoverage struct { Ratio *float64 `json:"ratio"` } type AdminSummaryRow struct { - ProjectID string `json:"project_id"` - KeyID *string `json:"key_id"` - AgentID *string `json:"agent_id"` - Assets *store.AdminAssetCounts `json:"assets"` - Sessions AdminSessionCounts `json:"sessions"` - Usage v1.TokenUsage `json:"usage"` - Coverage AdminUsageCoverage `json:"coverage"` - LastActiveAt *int64 `json:"last_active_at"` + ProjectID string `json:"project_id"` + KeyID *string `json:"key_id"` + AgentID *string `json:"agent_id"` + Assets *sessions.AdminAssetCounts `json:"assets"` + Sessions AdminSessionCounts `json:"sessions"` + Usage v1.TokenUsage `json:"usage"` + Coverage AdminUsageCoverage `json:"coverage"` + LastActiveAt *int64 `json:"last_active_at"` } type AdminSummaryResponse struct { Data []AdminSummaryRow `json:"data"` @@ -80,33 +79,33 @@ func (h *Handler) adminSummary(w http.ResponseWriter, r *http.Request) { group = "project" } if group != "project" && group != "agent" && group != "key" { - writeStoreError(w, r, sessions.ErrInvalidInput) + writeSessionsError(w, r, sessions.ErrInvalidInput) return } after, err := adminSummaryTime(r, "created_after") if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } before, err := adminSummaryTime(r, "created_before") if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } if after != nil && before != nil && !after.Before(*before) { - writeStoreError(w, r, sessions.ErrInvalidInput) + writeSessionsError(w, r, sessions.ErrInvalidInput) return } ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second) defer cancel() if group == "agent" && r.URL.Query().Get("project_id") == "" { - writeStoreError(w, r, sessions.ErrInvalidInput) + writeSessionsError(w, r, sessions.ErrInvalidInput) return } var page projects.Page if projectID := r.URL.Query().Get("project_id"); projectID != "" { if options.after != "" { - writeStoreError(w, r, sessions.ErrInvalidInput) + writeSessionsError(w, r, sessions.ErrInvalidInput) return } var binding projects.Binding @@ -125,7 +124,7 @@ func (h *Handler) adminSummary(w http.ResponseWriter, r *http.Request) { if group == "project" { groups[""] = &AdminSummaryRow{ProjectID: project.ID} } - counts, err := h.Admin.ReadAdminSummary(ctx, project.TenantID, store.AdminSummaryFilter{CreatedAfter: after, CreatedBefore: before}, func(session sessions.Session, creationKeyID *string) error { + counts, err := h.Admin.ReadAdminSummary(ctx, project.TenantID, sessions.AdminSummaryFilter{CreatedAfter: after, CreatedBefore: before}, func(session sessions.Session, creationKeyID *string) error { projected, err := sessionResponse(session, h.executorURL()) if err != nil { return err @@ -174,7 +173,7 @@ func (h *Handler) adminSummary(w http.ResponseWriter, r *http.Request) { return nil }) if err != nil { - writeStoreError(w, r, err) + writeSessionsError(w, r, err) return } if group == "project" { diff --git a/services/core/internal/api/dependencies.go b/services/core/internal/api/dependencies.go index 9c7a5e090..f4e2b408d 100644 --- a/services/core/internal/api/dependencies.go +++ b/services/core/internal/api/dependencies.go @@ -56,7 +56,6 @@ type Dependencies struct { EnvironmentsReader EnvironmentsReader ExecutorConnections ExecutorConnections Admin Admin - AdminRuntimeTargets AdminRuntimeTargets AdminAudit AdminAudit WriteAudit WriteAudit Metrics Metrics @@ -143,7 +142,7 @@ func (d Dependencies) validate() error { field{"Artifacts", d.Artifacts}, field{"ArtifactsReader", d.ArtifactsReader}, field{"SessionAdmin", d.SessionAdmin}, field{"Environments", d.Environments}, field{"EnvironmentsReader", d.EnvironmentsReader}, - field{"ExecutorConnections", d.ExecutorConnections}, field{"Admin", d.Admin}, field{"AdminRuntimeTargets", d.AdminRuntimeTargets}, field{"AdminAudit", d.AdminAudit}, field{"WriteAudit", d.WriteAudit}, + field{"ExecutorConnections", d.ExecutorConnections}, field{"Admin", d.Admin}, field{"AdminAudit", d.AdminAudit}, field{"WriteAudit", d.WriteAudit}, field{"Metrics", d.Metrics}, field{"RuntimeObservations", d.RuntimeObservations}, field{"RuntimeHistory", d.RuntimeHistory}, ); err != nil { return err diff --git a/services/core/internal/api/dependencies_test.go b/services/core/internal/api/dependencies_test.go index 38ce94de5..c7a5bf73b 100644 --- a/services/core/internal/api/dependencies_test.go +++ b/services/core/internal/api/dependencies_test.go @@ -42,7 +42,6 @@ type testFakes struct { environmentsReader *fakeEnvironmentsReader executorConnections *fakeExecutorConnections admin *fakeAdmin - adminRuntimeTargets *fakeAdminRuntimeTargets adminAudit *fakeAdminAudit writeAudit *fakeWriteAudit metrics *fakeMetrics @@ -87,7 +86,7 @@ func testDependencies(t testing.TB) (Dependencies, *testFakes) { artifacts: &fakeArtifacts{t: t}, artifactsReader: &fakeArtifactsReader{t: t}, sessionAdmin: &fakeSessionAdmin{t: t}, environments: &fakeEnvironments{t: t}, environmentsReader: &fakeEnvironmentsReader{t: t}, executorConnections: &fakeExecutorConnections{t: t}, - admin: &fakeAdmin{t: t}, adminRuntimeTargets: &fakeAdminRuntimeTargets{t: t}, adminAudit: &fakeAdminAudit{t: t}, writeAudit: &fakeWriteAudit{t: t}, metrics: &fakeMetrics{t: t}, + admin: &fakeAdmin{t: t}, adminAudit: &fakeAdminAudit{t: t}, writeAudit: &fakeWriteAudit{t: t}, metrics: &fakeMetrics{t: t}, runtimeObservations: &fakeRuntimeObservations{t: t}, runtimeHistory: &fakeRuntimeHistory{t: t}, installationBindings: &fakeInstallationBindings{t: t}, sessionAdmission: &fakeSessionAdmission{t: t}, inputAdmission: &fakeInputAdmission{t: t}, @@ -115,7 +114,7 @@ func testDependencies(t testing.TB) (Dependencies, *testFakes) { Artifacts: f.artifacts, ArtifactsReader: f.artifactsReader, SessionAdmin: f.sessionAdmin, - Environments: f.environments, EnvironmentsReader: f.environmentsReader, ExecutorConnections: f.executorConnections, Admin: f.admin, AdminRuntimeTargets: f.adminRuntimeTargets, AdminAudit: f.adminAudit, WriteAudit: f.writeAudit, + Environments: f.environments, EnvironmentsReader: f.environmentsReader, ExecutorConnections: f.executorConnections, Admin: f.admin, AdminAudit: f.adminAudit, WriteAudit: f.writeAudit, Metrics: f.metrics, RuntimeObservations: f.runtimeObservations, RuntimeHistory: f.runtimeHistory, }, f } diff --git a/services/core/internal/api/fakes_test.go b/services/core/internal/api/fakes_test.go index cf1565aee..7737d87bd 100644 --- a/services/core/internal/api/fakes_test.go +++ b/services/core/internal/api/fakes_test.go @@ -23,7 +23,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/skills" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/vaults" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/writeaudit" "github.com/google/uuid" @@ -42,23 +41,19 @@ func unexpectedCall(t testing.TB, method string) { } type fakeAdmin struct { - t testing.TB - readAdminSummary func(context.Context, string, store.AdminSummaryFilter, func(sessions.Session, *string) error) (store.AdminAssetCounts, error) + t testing.TB + readAdminSummary func(context.Context, string, sessions.AdminSummaryFilter, func(sessions.Session, *string) error) (sessions.AdminAssetCounts, error) + listAdminRuntimeTargets func(context.Context, []string, string, int, bool) (sessions.AdminRuntimeTargetPage, error) } -func (f *fakeAdmin) ReadAdminSummary(a0 context.Context, a1 string, a2 store.AdminSummaryFilter, a3 func(sessions.Session, *string) error) (store.AdminAssetCounts, error) { +func (f *fakeAdmin) ReadAdminSummary(a0 context.Context, a1 string, a2 sessions.AdminSummaryFilter, a3 func(sessions.Session, *string) error) (sessions.AdminAssetCounts, error) { if f.readAdminSummary == nil { unexpectedCall(f.t, "ReadAdminSummary") } return f.readAdminSummary(a0, a1, a2, a3) } -type fakeAdminRuntimeTargets struct { - t testing.TB - listAdminRuntimeTargets func(context.Context, []string, string, int, bool) (sessions.AdminRuntimeTargetPage, error) -} - -func (f *fakeAdminRuntimeTargets) ListAdminRuntimeTargets(a0 context.Context, a1 []string, a2 string, a3 int, a4 bool) (sessions.AdminRuntimeTargetPage, error) { +func (f *fakeAdmin) ListAdminRuntimeTargets(a0 context.Context, a1 []string, a2 string, a3 int, a4 bool) (sessions.AdminRuntimeTargetPage, error) { if f.listAdminRuntimeTargets == nil { unexpectedCall(f.t, "ListAdminRuntimeTargets") } diff --git a/services/core/internal/persistence/postgres/sessionpg/admin.go b/services/core/internal/persistence/postgres/sessionpg/admin.go index 3f49c20be..e23f728f7 100644 --- a/services/core/internal/persistence/postgres/sessionpg/admin.go +++ b/services/core/internal/persistence/postgres/sessionpg/admin.go @@ -3,6 +3,7 @@ package sessionpg import ( "context" "errors" + "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -53,3 +54,54 @@ func (s *Store) ListAdminRuntimeTargets(ctx context.Context, tenantIDs []string, } return page, nil } + +func (s *Store) ReadAdminSummary(ctx context.Context, tenantID string, filter sessions.AdminSummaryFilter, visit func(sessions.Session, *string) error) (sessions.AdminAssetCounts, error) { + var counts sessions.AdminAssetCounts + tenant, err := parseID(tenantID) + if err != nil { + return counts, err + } + err = s.units.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + raw, err := q.AdminAssetCounts(ctx, tenant) + if err != nil { + return err + } + counts = sessions.AdminAssetCounts(raw) + params := sqlc.AdminSummarySessionsParams{TenantID: tenant, CreatedAfter: timestamp(filter.CreatedAfter), CreatedBefore: timestamp(filter.CreatedBefore), AfterID: pgtype.UUID{Valid: true}} + for { + rows, err := q.AdminSummarySessions(ctx, params) + if err != nil { + return err + } + for _, row := range rows { + session, err := SessionFromRow(row.Session) + if err != nil { + return err + } + if session, err = LoadSessionActivity(ctx, q, session); err != nil { + return err + } + var creator *string + if row.CreationKeyID.Valid { + creator = &row.CreationKeyID.String + } + if err := visit(session, creator); err != nil { + return err + } + params.AfterID = row.Session.ID + } + if len(rows) < 100 { + return nil + } + } + }) + return counts, err +} + +func timestamp(value *time.Time) pgtype.Timestamptz { + if value == nil { + return pgtype.Timestamptz{} + } + return pgtype.Timestamptz{Time: *value, Valid: true} +} diff --git a/services/core/internal/persistence/postgres/sessionpg/admin_test.go b/services/core/internal/persistence/postgres/sessionpg/admin_test.go index 9d1e34609..43d37f264 100644 --- a/services/core/internal/persistence/postgres/sessionpg/admin_test.go +++ b/services/core/internal/persistence/postgres/sessionpg/admin_test.go @@ -61,3 +61,47 @@ func TestListAdminRuntimeTargetsPagesLiveSessionsOfTenants(t *testing.T) { t.Fatalf("malformed tenant: %v", err) } } + +func TestReadAdminSummaryVisitsTheSelectedSessions(t *testing.T) { + pool := pgtest.Open(t) + store := New(pgunit.NewPool(pool), nil) + tenant, created := uuid.NewString(), time.Now().UTC().Truncate(time.Second).Add(-time.Hour) + session := func(offset time.Duration, deleted bool) string { + id := uuid.NewString() + exec(t, pool, `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash, created_at, deleted_at) + VALUES ($1, $2, 'codex', $5, 'hash', $3, CASE WHEN $4::boolean THEN clock_timestamp() END)`, id, tenant, created.Add(offset), deleted, id) + return id + } + session(0, false) + turned, keyed := session(time.Second, false), session(2*time.Second, false) + session(2*time.Second, true) + session(3*time.Second, false) + turn := addTurn(t, pool, pgID(uuid.MustParse(turned)), "completed") + operation := uuid.NewString() + exec(t, pool, `INSERT INTO write_audit_operations(id, tenant_id, key_id, key_name, key_prefix, key_kind, action, resource_type, resource_id, request_id, trace_id) + VALUES ($1, $2, 'key', 'Key', 'sk', 'static', 'create', 'session', $3, $4, 'trace')`, operation, tenant, keyed, operation) + exec(t, pool, `INSERT INTO write_audit_owners(tenant_id, resource_type, resource_id, operation_id) VALUES ($1, 'session', $2, $3)`, tenant, keyed, operation) + + after, before := created.Add(time.Second), created.Add(3*time.Second) + visited := map[string]*string{} + counts, err := store.ReadAdminSummary(t.Context(), tenant, sessions.AdminSummaryFilter{CreatedAfter: &after, CreatedBefore: &before}, func(session sessions.Session, creator *string) error { + if session.ID == turned && (session.LastTurn == nil || session.LastTurn.ID != text(turn)) { + t.Errorf("activity of %s = %+v", turned, session.LastTurn) + } + visited[session.ID] = creator + return nil + }) + if err != nil || counts != (sessions.AdminAssetCounts{}) || len(visited) != 2 || visited[turned] != nil || visited[keyed] == nil || *visited[keyed] != "key" { + t.Fatalf("summary = %+v, %v, %v", counts, visited, err) + } + + exec(t, pool, `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash) + SELECT gen_random_uuid(), $1, 'codex', gen_random_uuid()::text, 'hash' FROM generate_series(1, 101)`, tenant) + total := 0 + if _, err := store.ReadAdminSummary(t.Context(), tenant, sessions.AdminSummaryFilter{}, func(sessions.Session, *string) error { total++; return nil }); err != nil || total != 105 { + t.Fatalf("unfiltered summary visited %d, %v", total, err) + } + if _, err := store.ReadAdminSummary(t.Context(), "not-a-tenant", sessions.AdminSummaryFilter{}, nil); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatalf("malformed tenant: %v", err) + } +} diff --git a/services/core/internal/sandbox/providers/configuration_flow_test.go b/services/core/internal/sandbox/providers/configuration_flow_test.go index 2fd9a314c..f400a64af 100644 --- a/services/core/internal/sandbox/providers/configuration_flow_test.go +++ b/services/core/internal/sandbox/providers/configuration_flow_test.go @@ -143,7 +143,7 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { Artifacts: struct{ api.Artifacts }{}, ArtifactsReader: struct{ api.ArtifactsReader }{}, SessionAdmin: struct{ api.SessionAdmin }{}, - Environments: struct{ api.Environments }{}, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, Admin: s, AdminRuntimeTargets: struct{ api.AdminRuntimeTargets }{}, AdminAudit: struct{ api.AdminAudit }{}, WriteAudit: struct{ api.WriteAudit }{}, + Environments: struct{ api.Environments }{}, EnvironmentsReader: struct{ api.EnvironmentsReader }{}, Admin: struct{ api.Admin }{}, 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/admin.go b/services/core/internal/sessions/admin.go index c68310671..3db0ce646 100644 --- a/services/core/internal/sessions/admin.go +++ b/services/core/internal/sessions/admin.go @@ -1,6 +1,9 @@ package sessions -import "context" +import ( + "context" + "time" +) // AdminRuntimeTarget names a live Session whose Runtime an administrator // observes, with its tenant. @@ -12,6 +15,20 @@ type AdminRuntimeTargetPage struct { HasMore bool } +// AdminSummaryFilter selects the Sessions created in [CreatedAfter, +// CreatedBefore); a nil bound is open. +type AdminSummaryFilter struct{ CreatedAfter, CreatedBefore *time.Time } + +// AdminAssetCounts counts a tenant's assets. +type AdminAssetCounts struct { + Agents int64 `json:"agents"` + Skills int64 `json:"skills"` + EnvironmentTemplates int64 `json:"environment_templates"` + Files int64 `json:"files"` + Vaults int64 `json:"vaults"` + Credentials int64 `json:"credentials"` +} + // AdminReader reads the administrator's cross-Project Session views. type AdminReader interface { // ListAdminRuntimeTargets pages the live Sessions of the tenants by @@ -20,4 +37,9 @@ type AdminReader interface { // ErrInvalidInput; an after Session that is malformed, deleted or // outside the tenants is ErrNotFound. ListAdminRuntimeTargets(ctx context.Context, tenants []string, after string, limit int, ascending bool) (AdminRuntimeTargetPage, error) + // ReadAdminSummary counts the tenant's assets and visits each live Session + // the filter selects, with GetSession's activity and the ID of the key that + // created it, nil when none is recorded, from one read-only snapshot. A + // malformed tenant is ErrInvalidInput; an error from visit ends the read. + ReadAdminSummary(ctx context.Context, tenant string, filter AdminSummaryFilter, visit func(Session, *string) error) (AdminAssetCounts, error) } diff --git a/services/core/internal/store/admin_summary.go b/services/core/internal/store/admin_summary.go deleted file mode 100644 index d7578aaf0..000000000 --- a/services/core/internal/store/admin_summary.go +++ /dev/null @@ -1,81 +0,0 @@ -package store - -import ( - "context" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -type AdminSummaryFilter struct{ CreatedAfter, CreatedBefore *time.Time } -type AdminAssetCounts struct { - Agents int64 `json:"agents"` - Skills int64 `json:"skills"` - EnvironmentTemplates int64 `json:"environment_templates"` - Files int64 `json:"files"` - Vaults int64 `json:"vaults"` - Credentials int64 `json:"credentials"` -} - -// ReadAdminSummary visits one space's Sessions from one read-only snapshot. The -// visitor reuses the API's Session projection instead of creating another status -// or usage model. Paging keeps the stored configurations out of an unbounded slice. -func (s *Store) ReadAdminSummary(ctx context.Context, tenantID string, filter AdminSummaryFilter, visit func(sessions.Session, *string) error) (AdminAssetCounts, error) { - var counts AdminAssetCounts - tenant, err := parseID(tenantID) - if err != nil { - return counts, err - } - if visit == nil || filter.CreatedAfter != nil && filter.CreatedBefore != nil && !filter.CreatedAfter.Before(*filter.CreatedBefore) { - return counts, sessions.ErrInvalidInput - } - err = s.pooled.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - raw, err := q.AdminAssetCounts(ctx, tenant) - if err != nil { - return err - } - counts = AdminAssetCounts{Agents: raw.Agents, Skills: raw.Skills, EnvironmentTemplates: raw.EnvironmentTemplates, Files: raw.Files, Vaults: raw.Vaults, Credentials: raw.Credentials} - params := sqlc.AdminSummarySessionsParams{TenantID: tenant, CreatedAfter: summaryTimestamp(filter.CreatedAfter), CreatedBefore: summaryTimestamp(filter.CreatedBefore), AfterID: pgtype.UUID{Valid: true}} - for { - rows, err := q.AdminSummarySessions(ctx, params) - if err != nil { - return err - } - for _, row := range rows { - session, err := sessionpg.SessionFromRow(row.Session) - if err != nil { - return err - } - session, err = sessionpg.LoadSessionActivity(ctx, q, session) - if err != nil { - return err - } - var creator *string - if row.CreationKeyID.Valid { - id := row.CreationKeyID.String - creator = &id - } - if err := visit(session, creator); err != nil { - return err - } - params.AfterID = row.Session.ID - } - if len(rows) < 100 { - return nil - } - } - }) - return counts, err -} - -func summaryTimestamp(value *time.Time) pgtype.Timestamptz { - if value == nil { - return pgtype.Timestamptz{} - } - return pgtype.Timestamptz{Time: *value, Valid: true} -} diff --git a/services/core/internal/store/public_handler_fixture_test.go b/services/core/internal/store/public_handler_fixture_test.go index 7db484e46..198938cbd 100644 --- a/services/core/internal/store/public_handler_fixture_test.go +++ b/services/core/internal/store/public_handler_fixture_test.go @@ -94,7 +94,7 @@ func publicHandler(t testing.TB, s *store.Store, db fixtureDB, keys fixtureKeyRe Subagents: sessionStore, Artifacts: sessionService, ArtifactsReader: sessionStore, - SessionAdmin: sessionStore, Environments: sessionService, EnvironmentsReader: sessionStore, Admin: s, AdminRuntimeTargets: sessionStore, AdminAudit: audit, WriteAudit: audit, + SessionAdmin: sessionStore, Environments: sessionService, EnvironmentsReader: sessionStore, Admin: sessionStore, AdminAudit: audit, WriteAudit: audit, ExecutorConnections: strict, Metrics: strict, RuntimeObservations: strict, RuntimeHistory: strict, } for _, c := range configure { From df2abab637d1dbbad36770a71214d40fcc83da8f Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 03:11:33 +0800 Subject: [PATCH 6/6] Move input admission and Environment input reservations into sessions (#441) * Move input admission and Environment input reservations into sessions * Admit input batches through one procedure and tighten reservation reads AdmitInputs admits a validated batch for every admission path, and admitInput is no longer exported. GetEnvironmentInputReservation reads in one snapshot through the bound Session transaction, the receipt and reservation loaders lose their forwarding helpers, FindInputReservation loads receipts only for a matching batch, and the pooled expiry calls in execution are bounded. TestTurnInputRetriesAndRestart again checks that admission leaves the stored Session unchanged and that its latest Turn survives a restart. --- services/core/IMPLEMENTATION.md | 6 +- .../native_classification_integration_test.go | 5 +- .../session_diagnostics_public_compat_test.go | 14 + .../internal/api/session_diagnostics_test.go | 5 +- .../archive_cancellation_cleanup_test.go | 4 +- services/core/internal/execution/delivery.go | 2 +- .../deployment_provider_observations_test.go | 4 +- .../execution/environment_admission.go | 6 +- .../core/internal/execution/message_input.go | 2 +- .../core/internal/execution/preparation.go | 4 +- .../internal/execution/prepared_dispatch.go | 10 +- services/core/internal/execution/worker.go | 2 +- .../internal/execution/worker_schedule.go | 8 +- .../core/internal/execution/worker_wakeup.go | 2 +- .../sessionpg/environment_inputs_test.go | 165 ++++ .../postgres/sessionpg/execution_inputs.go | 37 + .../sessionpg/execution_inputs_test.go | 776 ++++++++++++++++++ .../persistence/postgres/sessionpg/inputs.go | 247 ++++++ .../postgres/sessionpg/inputs_test.go | 521 ++++++++++++ .../postgres/sessionpg/turn_reads.go | 6 +- .../providers/configuration_flow_test.go | 2 +- .../core/internal/sessions/artifacts_test.go | 5 + .../core/internal/sessions/devices_test.go | 2 + services/core/internal/sessions/doc.go | 17 +- .../core/internal/sessions/environment.go | 2 + .../internal/sessions/environment_inputs.go | 159 ++++ .../sessions/environment_inputs_test.go | 165 ++++ services/core/internal/sessions/execution.go | 1 + .../sessions/execution_functions_test.go | 8 +- .../internal/sessions/execution_inputs.go | 103 +++ .../sessions/execution_inputs_test.go | 133 +++ services/core/internal/sessions/inputs.go | 332 ++++++++ .../core/internal/sessions/inputs_test.go | 322 ++++++++ services/core/internal/sessions/reader.go | 1 + services/core/internal/sessions/service.go | 1 + .../store/admin_session_archive_race_test.go | 2 +- .../store/admin_session_archive_test.go | 2 +- .../admin_session_archive_worker_http_test.go | 2 +- .../store/archive_cancellation_test.go | 4 +- .../internal/store/claude_execution_test.go | 2 +- .../internal/store/command_output_test.go | 4 +- .../creation_stream_settlement_public_test.go | 4 +- .../deployment_model_providers_http_test.go | 4 +- services/core/internal/store/dispatch_test.go | 4 +- .../store/environment_admission_test.go | 6 +- .../store/environment_claim_worker_test.go | 5 +- .../internal/store/environment_device_test.go | 2 +- .../environment_directory_active_test.go | 3 +- .../store/environment_expiry_worker_test.go | 6 +- .../store/environment_file_writes_test.go | 10 +- .../store/environment_initial_input_test.go | 16 +- .../store/environment_initial_public_test.go | 3 +- .../store/environment_input_activity.go | 23 - .../store/environment_input_activity_test.go | 38 +- .../store/environment_input_claim_test.go | 96 --- .../store/environment_input_expiry.go | 43 - .../store/environment_input_expiry_test.go | 103 +-- .../store/environment_input_migration_test.go | 36 - .../environment_input_settlement_test.go | 222 +---- .../core/internal/store/environment_inputs.go | 287 ------- .../internal/store/environment_inputs_test.go | 274 ------- .../store/environment_steering_order_test.go | 104 --- .../internal/store/environment_work_test.go | 8 +- .../store/environment_worker_helpers_test.go | 4 +- .../store/environment_worker_scan_test.go | 2 +- .../internal/store/environment_worker_test.go | 4 +- services/core/internal/store/execution.go | 18 +- .../internal/store/execution_events_test.go | 2 +- .../internal/store/execution_messages_test.go | 2 +- .../core/internal/store/execution_test.go | 28 +- .../internal/store/execution_tools_test.go | 2 +- services/core/internal/store/export_test.go | 23 +- .../store/function_execution_native_test.go | 2 +- .../internal/store/function_execution_test.go | 2 +- .../store/function_images_native_test.go | 2 +- .../store/function_input_execution_test.go | 4 +- .../store/function_inputs_public_test.go | 8 +- .../internal/store/function_inputs_test.go | 30 +- .../store/function_state_public_test.go | 2 +- .../internal/store/function_state_test.go | 2 +- .../internal/store/harness_onboarding_test.go | 2 +- ...sted_initialization_failure_public_test.go | 2 +- .../core/internal/store/input_batches_test.go | 189 ----- .../store/input_conflicts_public_test.go | 2 +- services/core/internal/store/inputs_test.go | 128 +++ .../core/internal/store/item_order_test.go | 6 +- .../core/internal/store/item_reads_test.go | 8 +- .../internal/store/list_cursor_public_test.go | 4 +- .../store/local_artifact_export_test.go | 3 +- .../local_environment_file_write_test.go | 4 +- .../store/local_environment_worker_test.go | 4 +- .../store/mcode_public_native_test.go | 2 +- .../store/message_images_native_test.go | 2 +- .../store/model_protocol_native_test.go | 2 +- .../store/prepared_dispatch_failure_test.go | 8 +- .../internal/store/prepared_dispatch_test.go | 4 +- .../internal/store/public_execution_test.go | 4 +- .../store/public_handler_fixture_test.go | 6 +- .../runtime_environment_terminal_test.go | 4 +- .../store/runtime_input_admission_test.go | 2 +- .../runtime_suspension_concurrency_test.go | 6 +- .../runtime_wake_hint_integration_test.go | 5 +- .../store/runtime_worker_recovery_test.go | 6 +- .../store/sandbox_deployment_switch_test.go | 2 +- services/core/internal/store/scheduling.go | 26 - .../store/self_hosted_cancel_public_test.go | 4 +- .../store/self_hosted_initial_public_test.go | 2 +- .../store/session_artifacts_public_test.go | 2 +- .../internal/store/session_artifacts_test.go | 2 +- .../store/session_creation_stream_test.go | 4 +- .../store/session_deletion_execution_test.go | 2 +- .../session_deletion_lifecycle_public_test.go | 11 +- .../internal/store/session_deletion_test.go | 20 +- .../session_environment_snapshot_test.go | 2 +- .../internal/store/session_events_test.go | 8 +- .../internal/store/session_initial_input.go | 15 +- .../store/session_initial_input_test.go | 6 +- .../internal/store/session_metadata_test.go | 2 +- .../store/session_write_audit_test.go | 8 +- services/core/internal/store/sessions.go | 2 +- .../internal/store/steering_receipts_test.go | 3 +- .../store/stream_authority_http_test.go | 2 +- .../store/token_usage_integration_test.go | 6 +- .../internal/store/tool_policy_native_test.go | 2 +- .../core/internal/store/turn_events_test.go | 2 +- services/core/internal/store/turn_inputs.go | 238 ------ .../core/internal/store/turn_inputs_test.go | 219 ----- services/core/internal/store/turns_test.go | 114 --- .../internal/store/worker_input_race_test.go | 2 +- .../store/worker_preparation_failure_test.go | 22 +- services/core/tests/fixtures/main.go | 11 +- 131 files changed, 3408 insertions(+), 2263 deletions(-) create mode 100644 services/core/internal/persistence/postgres/sessionpg/environment_inputs_test.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/execution_inputs.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/execution_inputs_test.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/inputs.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/inputs_test.go create mode 100644 services/core/internal/sessions/environment_inputs.go create mode 100644 services/core/internal/sessions/environment_inputs_test.go create mode 100644 services/core/internal/sessions/execution_inputs.go create mode 100644 services/core/internal/sessions/execution_inputs_test.go create mode 100644 services/core/internal/sessions/inputs_test.go delete mode 100644 services/core/internal/store/environment_input_activity.go delete mode 100644 services/core/internal/store/environment_input_claim_test.go delete mode 100644 services/core/internal/store/environment_input_expiry.go delete mode 100644 services/core/internal/store/environment_inputs.go delete mode 100644 services/core/internal/store/environment_inputs_test.go delete mode 100644 services/core/internal/store/environment_steering_order_test.go delete mode 100644 services/core/internal/store/input_batches_test.go create mode 100644 services/core/internal/store/inputs_test.go delete mode 100644 services/core/internal/store/scheduling.go delete mode 100644 services/core/internal/store/turn_inputs.go delete mode 100644 services/core/internal/store/turn_inputs_test.go delete mode 100644 services/core/internal/store/turns_test.go diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 3c5e0278d..9418835da 100644 --- a/services/core/IMPLEMENTATION.md +++ b/services/core/IMPLEMENTATION.md @@ -16,11 +16,11 @@ Two errors are shared across domains, each with one `api` helper: `textvalue.Err Shared vocabulary has one owner each, and domains use it rather than copy it. `internal/environmentconfig` owns Environment setup, Skills, Plugins and initial files with their validation and public metadata; `Setup.Validate` checks requested configuration, where a Skill may be an unresolved reference, and `Setup.ValidateInstalled` checks frozen, installable configuration. `internal/skills` owns `ParseVersion`, the canonical positive decimal Skill version. `internal/metadata` owns the metadata rules: `Validate` for the pair, key and value limits and U+0000, `ValidateStorable` for U+0000 alone, and `Encode` with its 64 KiB bound. `internal/jsonobject` owns `Normalize`, the stable encoding of stored JSON objects that snapshots and retry identities compare. These packages import no persistence. -`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, `LoadRequiredActions` lists the calls a Turn waits on, and `TransitionTurn` moves a root Turn's status as its current status allows and settles a Turn that ends. `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 and Turn operation loads the Turn 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. +`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, `AdmitInputs` admits a validated input batch in order, each input steering the active Turn, starting a Turn, requesting a cancellation or joining a function result to its call's Turn, `AdmitFunctionResult` saves a submitted function result for its Turn's call, `LoadRequiredActions` lists the calls a Turn waits on, and `TransitionTurn` moves a root Turn's status as its current status allows and settles a Turn that ends. `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 and Turn operation loads the Turn 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` or `TransitionTurn`, 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; among them `LoadSessionActivity`, the one Session activity read, adds the activity projection to the Session that `store`'s Session creation returns. `tests/fixtures` is test tooling that composes procedures over a pooled `WithSession`, which is why `WithSession` and `BindSession` stay exported after store goes; no production code may do this. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; `placementpg`, which loads the placement facts and reserves placements, sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. -`store` is transitional. `store.New` builds a pooled Store, and `store.NewExecution(s, lease)` builds the execution writer on a lease it borrows. An execution-only operation on a pooled Store fails with `store.ErrExecutionAuthority`. New adapters do not copy that check: their execution repositories require a `*pgunit.Lease` at construction, their public repositories expose no execution operation, and the check goes away with `store`. +`store` is transitional. `store.New` builds a pooled Store, and `store.NewExecution(s, lease)` builds the execution writer on a lease it borrows. Adapters' execution repositories require a `*pgunit.Lease` at construction, and their public repositories expose no execution operation. `cmd/server` owns the execution lease. It acquires one `pgunit.Lease`, builds every lease-bound adapter on it, and passes the lease and those adapters together as one `execution.Owner` to `execution.StartWorker`, which binds the Dispatcher to the Owner's execution operations through `Dispatcher.Bind`. If anything fails before that call, `cmd/server` closes the lease. From that call the Worker owns cleanup: a failed start closes the lease before it returns, and a started Worker closes it after `Run` has cancelled and drained its work. Each close runs under its own bounded deadline, independent of the cancelled request or run. Lease-bound adapters and the store writer borrow the lease and never close it, and the Worker uses the lease only through `Owner.Lease`, never through an adapter. Store integration tests start the Worker the same way through `startWorker`, and bind a Dispatcher that runs Turns without a Worker through `Dispatcher.Bind`. @@ -36,7 +36,7 @@ Domain owners, each with its PostgreSQL adapter under `internal/persistence/post - `deployment` and its subpackage `deployment/placement` (`deploymentpg`, `placementpg`): the sandbox deployment and its nodes: provider configuration and the sealed credential, the specification and retained generations, setup, update and switch, node enrollment, identity and authentication, generation configuration, capacity, presence, status, host history, reset, and the counts of nodes and sandboxes bound to the public address; and hosted runtime allocations: reservation with the dedicated device, compute ownership and settlement, observation diagnostics, activity, compute phases, wake receipts and cleanup, with the reads that schedule, discover and authorize them. It interprets Sandbox Provider declarations through the `providers.Registry` it is given, whose lookups return typed errors. `cmd/server` builds that registry and calls it only to build a direct Provider and to discover a Provider's configuration; it takes each setup's mode and declared operations from the `Setup` that `deployment` returns. `deployment/placement` owns hosted admission and placement: `cmd/server` builds one `placement.Rules` from the registry and the public URL, both fixed while Core runs, and gives it to `deployment.Service` and to `store`; its pure decisions admit a hosted Session, choose its node, and admit an allocation's reserved node and a restore on it, and its errors keep one status, code and message through `writeStoreError` and `writeDeploymentError`. `placementpg` loads the facts those decisions read and applies them on the caller's transaction-bound queries; it has no Store or transaction runner. The only other reader is `store`: Session creation decides admission and placement with those rules and `placementpg` inside its own transaction. Deployment changes and allocation writes run through `deployment.ExecutionOperations` on `deploymentpg.NewExecution(lease, …)`, which the Worker receives as `execution.Owner.Deployment`. Each allocation write locks the owning Session, decides in `deployment`, applies in `deploymentpg` and prunes the Session's change journal in the same leased transaction; cleanup settles the Session through the `sessions` procedures on a `sessionpg.SessionTx` bound to that transaction. The Worker reads the deployment, prepares a selection's setup and records live-compute activity through the pooled `deployment.Service` in `execution.Dispatcher.Deployment`, and lists the Sessions a reset still has to archive, schedules node lifecycles and reads allocations through the pooled `deployment.Reader` in `execution.Dispatcher.DeploymentReader`; a pooled read carries no lease, and the leased write that follows it rechecks the allocation's owner. Node management and reads use the pooled `deploymentpg.Store`, which `cmd/server` also reads the owner epoch from and runs the host-history sampler on. Provider calls run outside transactions, and the final transaction rechecks the expected generation. Session archive crosses the Session and the deployment, so it is a `deployment.ExecutionOperations` operation: `ArchiveSession`, which the administrator's archive route calls through `api.Execution.SessionArchive`, and `ArchiveResetSession`, which a reset calls, lock the Session through `deploymentpg`'s `WithSessionArchive`, check the deployment's generation and the running reset in `deployment`, then expire the Environment and cancel its work through the `sessions` procedures, revoke the allocation's device and request its cleanup, and record the audit entry in one leased transaction. `deployment.ObservationResolver` resolves a Session's Runtime observation target for `runtimeobs` from `sessions.SessionReader` and `deployment.Reader`. - `coremetrics` (`coremetricspg`): the Core metrics PostgreSQL holds: the root Turn queue counts, the root Turn history, read from one read-only snapshot, and the database size. `cmd/server`'s Core metrics source adds them to the process, pool, Worker and daemon registry measurements. - `runtimehistory` (`runtimehistory/postgresreader`, outside this directory): Runtime history samples: the periodic export, scoped reads and retention, which also prunes node-host samples. `runtimehistory.Service` scopes a read to the Session's hosted Environment through `sessions.EnvironmentReader`. -- `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 Sessions (their reads, the change journal and stream snapshot, diagnostics, the frozen execution configuration, measured usage and archive state, metadata updates, deletion and the public write audit), Turns (their reads and the execution work scan), the model provider a Session froze, which `sessionpg.Store` opens with the credential key, 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, which include a Session's execution binding and the execution device list, their creation, revocation, heartbeats and Runtime enrollment, and the archived cancellation receipt the daemon gateway reads), the administrator's views across Projects (a Project's asset counts and Sessions for the summary, and the Sessions whose Runtime the administrator observes), 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. Turn transitions, execution completion, which alone publishes or discards a Turn's staged Artifacts, the start of Artifact capture, 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`. +- `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 Sessions (their reads, the change journal and stream snapshot, diagnostics, the frozen execution configuration, measured usage and archive state, metadata updates, deletion and the public write audit), Turns (their reads and the execution work scan), input admission (a public input batch, Environment input reservation and the expiry of one reservation, with the reads of a Turn's admitted inputs, a reservation and the Environment input work scan), the model provider a Session froze, which `sessionpg.Store` opens with the credential key, 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, which include a Session's execution binding and the execution device list, their creation, revocation, heartbeats and Runtime enrollment, and the archived cancellation receipt the daemon gateway reads), the administrator's views across Projects (a Project's asset counts and Sessions for the summary, and the Sessions whose Runtime the administrator observes), 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. Turn transitions, execution completion, which alone publishes or discards a Turn's staged Artifacts, the start of Artifact capture, function calls and their application receipts, the Turn execution journal, Environment initialization, connection observations and their reconciliation, Session device binding, file-write reservation and settlement, and the promotion, failure and bulk expiry of Environment input reservations run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. ## Request handling diff --git a/services/core/internal/api/native_classification_integration_test.go b/services/core/internal/api/native_classification_integration_test.go index 9ce5ea206..12dc3754f 100644 --- a/services/core/internal/api/native_classification_integration_test.go +++ b/services/core/internal/api/native_classification_integration_test.go @@ -25,10 +25,7 @@ func TestNativeClassificationPostgresRoundTripAndPublicPrivacy(t *testing.T) { if err != nil { t.Fatal(err) } - receipt, err := s.SubmitMessage(t.Context(), tenant, session.ID, "input", json.RawMessage(`{"text":"test"}`)) - if err != nil { - t.Fatal(err) - } + receipt := submitMessage(t, pool, tenant, session.ID, "input", json.RawMessage(`{"text":"test"}`)) transitionTurn(t, pool, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) status := 503 result := execution.Result{ErrorCode: "engine_failed", Error: "Bearer secret-canary https://private.example/key", EngineErrorCode: code, EngineHTTPStatus: &status, Done: proto.DonePayload{Usage: proto.Usage{InputTokens: 7, OutputTokens: 3}, Metadata: map[string]any{proto.DoneMetaAgentSessionID: "native-secret-canary"}}} diff --git a/services/core/internal/api/session_diagnostics_public_compat_test.go b/services/core/internal/api/session_diagnostics_public_compat_test.go index c555210b3..f6b56f881 100644 --- a/services/core/internal/api/session_diagnostics_public_compat_test.go +++ b/services/core/internal/api/session_diagnostics_public_compat_test.go @@ -73,6 +73,20 @@ func databaseSessionReads(pool *pgxpool.Pool) func(*Dependencies, *testFakes) { } } +// submitMessage admits one message input through the Session service on pool. +func submitMessage(t *testing.T, pool *pgxpool.Pool, tenant, session, key string, payload json.RawMessage) sessions.InputReceipt { + t.Helper() + service, err := sessions.NewService(sessionpg.New(pgunit.NewPool(pool), nil)) + if err != nil { + t.Fatal(err) + } + receipts, err := service.SubmitInputs(t.Context(), tenant, session, key, []sessions.Input{{Kind: "message", Payload: payload}}) + if err != nil { + t.Fatal(err) + } + return receipts[0] +} + // transitionTurn moves the Turn as the execution owner does, over a pooled // Session transaction. func transitionTurn(t *testing.T, pool *pgxpool.Pool, tenant, session, turn string, transition sessions.TurnTransition) { diff --git a/services/core/internal/api/session_diagnostics_test.go b/services/core/internal/api/session_diagnostics_test.go index 448a959df..0470488c1 100644 --- a/services/core/internal/api/session_diagnostics_test.go +++ b/services/core/internal/api/session_diagnostics_test.go @@ -19,10 +19,7 @@ func TestDiagnosticsCoreHandlerDatabaseBoundary(t *testing.T) { if err != nil { t.Fatal(err) } - receipt, err := s.SubmitMessage(t.Context(), tenant, session.ID, "input", json.RawMessage(`{"text":"input-secret-canary"}`)) - if err != nil { - t.Fatal(err) - } + receipt := submitMessage(t, pool, tenant, session.ID, "input", json.RawMessage(`{"text":"input-secret-canary"}`)) transitionTurn(t, pool, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) transitionTurn(t, pool, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed, Outcome: json.RawMessage(`{"error_code":"device_disconnected","error":"Bearer raw-secret-canary https://private.example/key","done":{"native_id":"secret-native-canary"}}`)}) base := adminSessionsPath + session.ID diff --git a/services/core/internal/execution/archive_cancellation_cleanup_test.go b/services/core/internal/execution/archive_cancellation_cleanup_test.go index 822676a52..34ac69f46 100644 --- a/services/core/internal/execution/archive_cancellation_cleanup_test.go +++ b/services/core/internal/execution/archive_cancellation_cleanup_test.go @@ -92,10 +92,12 @@ func TestArchiveWaitingCleanupReceiptBarrier(t *testing.T) { if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(t.Context(), project.TenantID, session.ID, "start", json.RawMessage(`{"text":"run"}`)) + _, service := testSessions(t, pool, nil) + inputs, err := service.SubmitInputs(t.Context(), project.TenantID, session.ID, "start", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"run"}`)}}) if err != nil { t.Fatal(err) } + input := inputs[0] // This fixture isolates lifecycle ordering. Protocol-driven waiting is // independently exercised in TestArchiveWaitingCancellationReceipts. for _, transition := range []sessions.TurnTransition{{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}, {ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnWaiting}} { diff --git a/services/core/internal/execution/delivery.go b/services/core/internal/execution/delivery.go index ab20d7d7c..b18f4f86e 100644 --- a/services/core/internal/execution/delivery.go +++ b/services/core/internal/execution/delivery.go @@ -303,7 +303,7 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe continue } if pending == nil { - inputs, err := d.Store.ListTurnInputs(ctx, tenantID, sessionID, request.RunID, result.AppliedThrough, 1) + inputs, err := d.SessionsReader.ListTurnInputs(ctx, tenantID, sessionID, request.RunID, result.AppliedThrough, 1) if err != nil { result.ErrorCode = "execution_state_unavailable" if ctx.Err() != nil { diff --git a/services/core/internal/execution/deployment_provider_observations_test.go b/services/core/internal/execution/deployment_provider_observations_test.go index fa11d22b2..d8a45d3e4 100644 --- a/services/core/internal/execution/deployment_provider_observations_test.go +++ b/services/core/internal/execution/deployment_provider_observations_test.go @@ -82,10 +82,12 @@ func newFinishObservationFixture(t *testing.T, maxConnections int32) finishObser } func (f finishObservationFixture) start(t *testing.T) sessions.InputReceipt { t.Helper() - receipt, err := f.s.SubmitMessage(t.Context(), f.tenant, f.session.ID, uuid.NewString(), json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"fixture"}]}]}`)) + _, service := testSessions(t, f.pool, nil) + receipts, err := service.SubmitInputs(t.Context(), f.tenant, f.session.ID, uuid.NewString(), []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"fixture"}]}]}`)}}) if err != nil { t.Fatal(err) } + receipt := receipts[0] if _, err = f.execution.TransitionTurn(t.Context(), f.tenant, f.session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/environment_admission.go b/services/core/internal/execution/environment_admission.go index 835027d90..d078186e7 100644 --- a/services/core/internal/execution/environment_admission.go +++ b/services/core/internal/execution/environment_admission.go @@ -97,7 +97,7 @@ func (w *Worker) submitEnvironmentInputs(ctx context.Context, session sessions.S changed, unsubscribe := w.dispatcher.notifications.subscribe(session.TenantID, session.ID) defer unsubscribe() reserve, cancel := context.WithTimeout(ctx, 5*time.Second) - reservation, err := w.admission.ReserveEnvironmentInput(reserve, session.TenantID, session.ID, key, inputs) + reservation, err := w.dispatcher.Sessions.ReserveEnvironmentInput(reserve, session.TenantID, session.ID, key, inputs) cancel() if err != nil { return nil, err @@ -140,9 +140,9 @@ func (w *Worker) environmentInputOutcome(ctx context.Context, session sessions.S defer cancel() // The database rechecks its clock under the Session lock before settlement. if !time.Now().Before(reservation.Deadline) { - return w.admission.ExpireEnvironmentInput(read, session.TenantID, session.ID, reservation.ID) + return w.dispatcher.Sessions.ExpireEnvironmentInput(read, session.TenantID, session.ID, reservation.ID) } - return w.admission.GetEnvironmentInputReservation(read, session.TenantID, session.ID, reservation.ID) + return w.dispatcher.SessionsReader.GetEnvironmentInputReservation(read, session.TenantID, session.ID, reservation.ID) } func (w *Worker) checkAdmissionOwnership(ctx context.Context) error { diff --git a/services/core/internal/execution/message_input.go b/services/core/internal/execution/message_input.go index 5a44766d8..328567f6f 100644 --- a/services/core/internal/execution/message_input.go +++ b/services/core/internal/execution/message_input.go @@ -38,7 +38,7 @@ func messageInput(raw json.RawMessage) (proto.MessageInput, error) { } func (d *Dispatcher) initialInput(ctx context.Context, tenant, session, turn string) (proto.MessageInput, int64, error) { - inputs, err := d.Store.ListTurnInputs(ctx, tenant, session, turn, 0, 100) + inputs, err := d.SessionsReader.ListTurnInputs(ctx, tenant, session, turn, 0, 100) if err != nil { return nil, 0, err } diff --git a/services/core/internal/execution/preparation.go b/services/core/internal/execution/preparation.go index 59ea9bfb4..ca815d02e 100644 --- a/services/core/internal/execution/preparation.go +++ b/services/core/internal/execution/preparation.go @@ -97,7 +97,9 @@ func (d *Dispatcher) awaitPreparation(ctx context.Context, tenant, session strin case <-ctx.Done(): return pending, ctx.Err() case <-tick.C: - current, err := d.Store.ExpireEnvironmentInput(ctx, tenant, session, pending.ID) + expire, cancel := context.WithTimeout(ctx, 5*time.Second) + current, err := d.Sessions.ExpireEnvironmentInput(expire, tenant, session, pending.ID) + cancel() if err != nil || current.State != sessions.EnvironmentInputPending { return current, err } diff --git a/services/core/internal/execution/prepared_dispatch.go b/services/core/internal/execution/prepared_dispatch.go index 1584b7442..6cb5817f6 100644 --- a/services/core/internal/execution/prepared_dispatch.go +++ b/services/core/internal/execution/prepared_dispatch.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "strings" + "time" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" @@ -17,12 +18,15 @@ type EnvironmentRun struct { } // RunEnvironmentInput reserves a Turn on the Session-owned Runtime Executor. It -// checks lease, the lease d.Store was built on, before any Runtime preparation. +// checks lease, the lease the Dispatcher's execution operations hold, before +// any Runtime preparation. func (d *Dispatcher) RunEnvironmentInput(ctx context.Context, lease Ownership, tenantID, sessionID, reservationID string) (run EnvironmentRun, err error) { if err = lease.CheckOwnership(ctx); err != nil { return run, err } - run.Reservation, err = d.Store.ExpireEnvironmentInput(ctx, tenantID, sessionID, reservationID) + expire, cancel := context.WithTimeout(ctx, 5*time.Second) + run.Reservation, err = d.Sessions.ExpireEnvironmentInput(expire, tenantID, sessionID, reservationID) + cancel() if err != nil || run.Reservation.State != sessions.EnvironmentInputPending { return run, err } @@ -92,7 +96,7 @@ func (d *Dispatcher) RunEnvironmentInput(ctx context.Context, lease Ownership, t if err := d.messageInputSupport(peer, session.Engine, snapshot, messages); err != nil { return run, err } - promoted, err := d.Store.PromoteEnvironmentInput(owner, tenantID, sessionID, reservationID) + promoted, err := d.sessionExecution.PromoteEnvironmentInput(owner, tenantID, sessionID, reservationID) if errors.Is(err, sessions.ErrTurnConflict) { // A rejected claim leaves the reservation pending for a later attempt. return run, err diff --git a/services/core/internal/execution/worker.go b/services/core/internal/execution/worker.go index 0b2b48a0b..6ad0bcc06 100644 --- a/services/core/internal/execution/worker.go +++ b/services/core/internal/execution/worker.go @@ -303,7 +303,7 @@ func (w *Worker) Run(ctx context.Context) (runErr error) { return err } if maintenance { - if _, err := w.dispatcher.Store.ExpireEnvironmentInputs(ctx); err != nil { + if _, err := w.dispatcher.sessionExecution.ExpireEnvironmentInputs(ctx); err != nil { w.observeSchedulerPoll(0, err) return err } diff --git a/services/core/internal/execution/worker_schedule.go b/services/core/internal/execution/worker_schedule.go index fbc0098d8..de5ac027a 100644 --- a/services/core/internal/execution/worker_schedule.go +++ b/services/core/internal/execution/worker_schedule.go @@ -34,7 +34,7 @@ func (s *workerSchedule) selectWork(ctx context.Context, w *Worker, devices []st } var environments []sessions.EnvironmentInputWork if !time.Now().Before(s.nextEnvironmentScan) { - environments, err = w.dispatcher.Store.ListEnvironmentInputWork(ctx, s.environmentCursor, devices) + environments, err = w.dispatcher.SessionsReader.ListEnvironmentInputWork(ctx, s.environmentCursor, devices) if err != nil { return nil, err } @@ -43,7 +43,7 @@ func (s *workerSchedule) selectWork(ctx context.Context, w *Worker, devices []st s.environmentCursor = "" // Retry the first page now instead of spending a scan interval on EOF. // A single refill preserves the candidate bound and cannot spin when empty. - environments, err = w.dispatcher.Store.ListEnvironmentInputWork(ctx, "", devices) + environments, err = w.dispatcher.SessionsReader.ListEnvironmentInputWork(ctx, "", devices) if err != nil { return nil, err } @@ -94,10 +94,10 @@ func (w *Worker) runEnvironmentInput(ctx context.Context, item scheduledWork) er return nil } if errors.Is(err, ErrModelProviderRequired) { - return w.dispatcher.Store.FailEnvironmentInput(ctx, item.TenantID, item.SessionID, item.reservationID, "model_provider_required") + return w.dispatcher.sessionExecution.FailEnvironmentInput(ctx, item.TenantID, item.SessionID, item.reservationID, "model_provider_required") } if errors.Is(err, errPreparationFailed) && run.Reservation.State == sessions.EnvironmentInputPending { - return w.dispatcher.Store.FailEnvironmentInput(ctx, item.TenantID, item.SessionID, item.reservationID, "runtime_preparation_failed") + return w.dispatcher.sessionExecution.FailEnvironmentInput(ctx, item.TenantID, item.SessionID, item.reservationID, "runtime_preparation_failed") } if run.Reservation.State == sessions.EnvironmentInputAdmitted { return err diff --git a/services/core/internal/execution/worker_wakeup.go b/services/core/internal/execution/worker_wakeup.go index 403e2bead..153a6b962 100644 --- a/services/core/internal/execution/worker_wakeup.go +++ b/services/core/internal/execution/worker_wakeup.go @@ -16,7 +16,7 @@ func (w *Worker) wakeScheduler() { } func (w *Worker) admitInputs(ctx context.Context, tenant, session, key string, inputs []sessions.Input) ([]sessions.InputReceipt, error) { - receipts, err := w.admission.SubmitInputs(ctx, tenant, session, key, inputs) + receipts, err := w.dispatcher.Sessions.SubmitInputs(ctx, tenant, session, key, inputs) if err == nil { w.wakeScheduler() w.dispatcher.notifications.notify(tenant, session) diff --git a/services/core/internal/persistence/postgres/sessionpg/environment_inputs_test.go b/services/core/internal/persistence/postgres/sessionpg/environment_inputs_test.go new file mode 100644 index 000000000..9a000de16 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/environment_inputs_test.go @@ -0,0 +1,165 @@ +package sessionpg + +import ( + "context" + "encoding/json" + "errors" + "reflect" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var reservedBatch = []sessions.Input{messageInput("first"), messageInput("second")} + +// environmentInputHistory checks the Session's Turns and admitted inputs, one +// Item per input, and that no input left Turn or Item events. +func environmentInputHistory(t *testing.T, pool *pgxpool.Pool, session pgtype.UUID, turns, inputs int) { + t.Helper() + var gotTurns, gotInputs, items, events int + err := pool.QueryRow(t.Context(), ` + SELECT (SELECT count(*) FROM turns WHERE session_id=$1), + (SELECT count(*) FROM turn_inputs WHERE session_id=$1), + (SELECT count(*) FROM session_items WHERE session_id=$1), + (SELECT count(*) FROM session_events WHERE session_id=$1 + AND (payload ? 'turn' OR payload->'event' ? 'item'))`, session).Scan(&gotTurns, &gotInputs, &items, &events) + if err != nil || gotTurns != turns || gotInputs != inputs || items != inputs || (inputs == 0 && events != 0) { + t.Fatal("history", gotTurns, gotInputs, items, events, err) + } +} + +// reserve reserves the two-message batch under key. +func reserve(t *testing.T, service *sessions.Service, tenant, session pgtype.UUID, key string) sessions.EnvironmentInputReservation { + t.Helper() + got, err := service.ReserveEnvironmentInput(t.Context(), text(tenant), text(session), key, reservedBatch) + if err != nil { + t.Fatal(err) + } + return got +} + +// passDeadline moves the reservation's deadline into the past. +func passDeadline(t *testing.T, pool *pgxpool.Pool, reservation string) { + t.Helper() + exec(t, pool, `UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE id=$1`, reservation) +} + +// cancelPending cancels the Session's pending input as Session cancellation +// does, then reads the reservation back. +func cancelPending(ctx context.Context, pool *pgxpool.Pool, store *Store, tenant, session pgtype.UUID, reservation string) (sessions.EnvironmentInputReservation, error) { + err := WithSession(ctx, pgunit.NewPool(pool), tenant, session, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + if err := locked.Public(); err != nil { + return err + } + bound := BindSession(q, tenant, session) + return sessions.TrackInputActivity(ctx, bound, bound.CancelPendingInput) + }) + if err != nil { + return sessions.EnvironmentInputReservation{}, err + } + return store.GetEnvironmentInputReservation(ctx, text(tenant), text(session), reservation) +} + +// Concurrent equivalent reservations from two pools share one pending +// identity and deadline, which a restart keeps. +func TestEnvironmentInputReservationConcurrentIdentity(t *testing.T) { + pool := pgtest.Open(t) + _, service := stagingService(t, pool) + _, other := stagingService(t, pgtest.Open(t)) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + ctx := t.Context() + batch := []sessions.Input{ + {Kind: "message", Payload: json.RawMessage(`{"text":"first","detail":{"a":1,"b":2}}`)}, + messageInput("second"), + } + const count = 8 + results := make(chan sessions.EnvironmentInputReservation, count) + var wg sync.WaitGroup + for i := range count { + wg.Go(func() { + admission := service + inputs := append([]sessions.Input(nil), batch...) + if i%2 == 0 { + admission = other + inputs[0].Payload = json.RawMessage(` { "detail": {"b": 2, "a": 1}, "text": "first" } `) + } + got, err := admission.ReserveEnvironmentInput(ctx, text(tenant), text(session), "request", inputs) + if err != nil { + t.Error(err) + return + } + results <- got + }) + } + wg.Wait() + close(results) + var first sessions.EnvironmentInputReservation + received := 0 + for result := range results { + received++ + if first.ID == "" { + first = result + } + if !reflect.DeepEqual(first, result) { + t.Fatal("reservation identity changed", first, result) + } + } + if received != count || first.State != sessions.EnvironmentInputPending || first.ID == "" || first.Deadline.Sub(first.CreatedAt) != 5*time.Minute || first.SettledAt != nil || len(first.Receipts) != 0 { + t.Fatal("invalid pending result", received, first) + } + environmentInputHistory(t, pool, session, 0, 0) + for _, changed := range [][]sessions.Input{batch[:1], {batch[1], batch[0]}, {messageInput("changed"), batch[1]}} { + if _, err := service.ReserveEnvironmentInput(ctx, text(tenant), text(session), "request", changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { + t.Fatal("changed request accepted", err) + } + } + if _, err := other.ReserveEnvironmentInput(ctx, text(tenant), text(session), "other", batch); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatal("second pending request accepted", err) + } + pool.Close() + _, restarted := stagingService(t, pgtest.Open(t)) + got, err := restarted.ReserveEnvironmentInput(ctx, text(tenant), text(session), "request", batch) + if err != nil || !reflect.DeepEqual(first, got) { + t.Fatal("restart changed deadline or identity", got, err) + } +} + +// A key that direct admission already used keeps its receipts and gains no +// reservation, and its retry leaves newer pending input alone. +func TestEnvironmentInputReservationKeepsEarlierDirectIdentity(t *testing.T) { + pool := pgtest.Open(t) + store, service := stagingService(t, pool) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + ctx := t.Context() + input := messageInput("already admitted") + receipts, err := service.SubmitInputs(ctx, text(tenant), text(session), "direct", []sessions.Input{input}) + if err != nil { + t.Fatal(err) + } + move(t, pool, tenant, session, receipts[0].TurnID, sessions.TurnQueued, sessions.TurnInProgress) + move(t, pool, tenant, session, receipts[0].TurnID, sessions.TurnInProgress, sessions.TurnCompleted) + pending := reserve(t, service, tenant, session, "new") + got, err := service.ReserveEnvironmentInput(ctx, text(tenant), text(session), "direct", []sessions.Input{input}) + if err != nil || got.State != sessions.EnvironmentInputAdmitted || got.ID != "" || !got.Deadline.IsZero() || len(got.Receipts) != 1 || got.Receipts[0].Sequence != receipts[0].Sequence { + t.Fatal("direct admission gained a reservation", got, err) + } + if _, err := service.ReserveEnvironmentInput(ctx, text(tenant), text(session), "direct", []sessions.Input{messageInput("changed")}); !errors.Is(err, sessions.ErrIdempotencyConflict) { + t.Fatal(err) + } + retry, err := service.SubmitInputs(ctx, text(tenant), text(session), "direct", []sessions.Input{input}) + if err != nil || len(retry) != 1 || !retry[0].Replayed { + t.Fatal(retry, err) + } + retained, err := store.GetEnvironmentInputReservation(ctx, text(tenant), text(session), pending.ID) + if err != nil || !reflect.DeepEqual(retained, pending) { + t.Fatal("old retry affected new pending input", retained, err) + } +} diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_inputs.go b/services/core/internal/persistence/postgres/sessionpg/execution_inputs.go new file mode 100644 index 000000000..77c88f0a9 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution_inputs.go @@ -0,0 +1,37 @@ +package sessionpg + +import ( + "context" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +func (e *Execution) WithInputs(ctx context.Context, tenant, session string, apply func(context.Context, sessions.InputTx) error) error { + return withInputs(ctx, e.lease, tenant, session, apply) +} + +// WithDueInputReservations locks each due reservation's Session with the scan +// itself, skipping Sessions another transaction holds, and prunes each +// Session's change journal after apply. +func (e *Execution) WithDueInputReservations(ctx context.Context, apply func(context.Context, sessions.InputTx, string) error) error { + return e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + rows, err := q.ListDueEnvironmentInputs(ctx) + if err != nil { + return err + } + for _, row := range rows { + if err := apply(ctx, BindSession(q, row.TenantID, row.SessionID), uuid.UUID(row.ID.Bytes).String()); err != nil { + return err + } + if err := PruneChanges(ctx, q, row.SessionID); err != nil { + return err + } + } + return nil + }) +} diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_inputs_test.go b/services/core/internal/persistence/postgres/sessionpg/execution_inputs_test.go new file mode 100644 index 000000000..22807a559 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution_inputs_test.go @@ -0,0 +1,776 @@ +package sessionpg + +import ( + "context" + "errors" + "reflect" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// leasedInputs opens a database of its own and returns its pool, the Session +// store and service on the pool, and the execution operations on its lease. +func leasedInputs(t *testing.T) (*pgxpool.Pool, *Store, *sessions.Service, *sessions.ExecutionOperations) { + t.Helper() + pool := pgtest.OpenIsolated(t, nil) + store, service := stagingService(t, pool) + operations, _ := sessionExecution(t, pool) + return pool, store, service, operations +} + +// terminateLeaseOwner ends the backend that holds the execution lease of +// pool's database. +func terminateLeaseOwner(t *testing.T, pool *pgxpool.Pool) { + t.Helper() + var killed bool + err := pool.QueryRow(t.Context(), `SELECT pg_terminate_backend(pid, 1000) FROM pg_locks WHERE locktype='advisory' AND granted AND objsubid=1 + AND classid::bigint * 4294967296 + objid::bigint = 706172736172 + AND database=(SELECT oid FROM pg_database WHERE datname=current_database())`).Scan(&killed) + if err != nil || !killed { + t.Fatal("terminate execution lease owner", killed, err) + } +} + +func reservationState(t *testing.T, store *Store, tenant, session pgtype.UUID, reservation string) string { + t.Helper() + got, err := store.GetEnvironmentInputReservation(t.Context(), text(tenant), text(session), reservation) + if err != nil { + t.Fatal(err) + } + return got.State +} + +// lockSession holds the Session's lock in an open transaction and returns the +// holder's backend PID; the transaction rolls back when the test ends unless +// release commits it first. +func lockSession(ctx context.Context, t *testing.T, pool *pgxpool.Pool, session pgtype.UUID) (int32, func(sql string, args ...any)) { + t.Helper() + tx, err := pool.Begin(ctx) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = tx.Rollback(context.Background()) }) + var pid int32 + if err := tx.QueryRow(ctx, "SELECT pg_backend_pid() FROM sessions WHERE id=$1 FOR UPDATE", session).Scan(&pid); err != nil { + t.Fatal(err) + } + return pid, func(sql string, args ...any) { + t.Helper() + if sql != "" { + if _, err := tx.Exec(ctx, sql, args...); err != nil { + t.Fatal(err) + } + } + if err := tx.Commit(ctx); err != nil { + t.Fatal(err) + } + } +} + +// awaitBlocked waits until a backend waits on holder's lock and returns its +// PID. +func awaitBlocked(ctx context.Context, t *testing.T, pool *pgxpool.Pool, holder int32) int32 { + t.Helper() + for ctx.Err() == nil { + var blocked int32 + if err := pool.QueryRow(ctx, "SELECT pid FROM pg_stat_activity WHERE $1=ANY(pg_blocking_pids(pid)) ORDER BY pid LIMIT 1", holder).Scan(&blocked); err == nil { + return blocked + } + time.Sleep(5 * time.Millisecond) + } + t.Fatal("transaction did not wait for the Session lock") + return 0 +} + +// Direct input waits behind a pending reservation; promotion admits it once, +// and every retry, after a new execution owner and pool too, replays that +// admission. +func TestEnvironmentInputReservationPromotionAndDirectRetries(t *testing.T) { + pool := pgtest.OpenIsolated(t, nil) + store, service := stagingService(t, pool) + operations, lease := sessionExecution(t, pool) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + ctx := t.Context() + first := reserve(t, service, tenant, session, "pending") + for _, request := range []struct { + key string + inputs []sessions.Input + want error + }{ + {"pending", first.Inputs, sessions.ErrTurnConflict}, + {"pending", []sessions.Input{messageInput("changed")}, sessions.ErrIdempotencyConflict}, + {"later", []sessions.Input{messageInput("later")}, sessions.ErrTurnConflict}, + {"cancel", []sessions.Input{cancelInput}, sessions.ErrTurnConflict}, + } { + if _, err := service.SubmitInputs(ctx, text(tenant), text(session), request.key, request.inputs); !errors.Is(err, request.want) { + t.Fatal("direct path bypassed reservation", request.key, err) + } + } + environmentInputHistory(t, pool, session, 0, 0) + promoted, err := operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), first.ID) + if err != nil || promoted.State != sessions.EnvironmentInputAdmitted || promoted.SettledAt == nil || len(promoted.Receipts) != 2 || !promoted.Deadline.Equal(first.Deadline) { + t.Fatal(promoted, err) + } + for i, receipt := range promoted.Receipts { + if receipt.Replayed || receipt.TurnID == "" || receipt.TurnID != promoted.Receipts[0].TurnID || (i > 0 && receipt.Sequence <= promoted.Receipts[i-1].Sequence) { + t.Fatal("promotion receipts", promoted.Receipts) + } + } + environmentInputHistory(t, pool, session, 1, 2) + for _, read := range []func() (sessions.EnvironmentInputReservation, error){ + func() (sessions.EnvironmentInputReservation, error) { + return operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), first.ID) + }, + func() (sessions.EnvironmentInputReservation, error) { + return store.GetEnvironmentInputReservation(ctx, text(tenant), text(session), first.ID) + }, + func() (sessions.EnvironmentInputReservation, error) { + return service.ReserveEnvironmentInput(ctx, text(tenant), text(session), "pending", first.Inputs) + }, + } { + retry, err := read() + if err != nil || retry.ID != first.ID || !retry.Deadline.Equal(first.Deadline) || retry.State != sessions.EnvironmentInputAdmitted || len(retry.Receipts) != 2 { + t.Fatal(retry, err) + } + for i, receipt := range retry.Receipts { + if !receipt.Replayed || receipt.Sequence != promoted.Receipts[i].Sequence || receipt.TurnID != promoted.Receipts[i].TurnID { + t.Fatal("retry changed admission", receipt) + } + } + } + retry, err := service.SubmitInputs(ctx, text(tenant), text(session), "pending", first.Inputs) + if err != nil || len(retry) != 2 || !retry[0].Replayed || retry[0].Sequence != promoted.Receipts[0].Sequence { + t.Fatal("direct retry after promotion", retry, err) + } + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, pool) + if err := lease.Close(ctx); err != nil { + t.Fatal(err) + } + awaitRelease() + restarted, err := pgxpool.NewWithConfig(ctx, pool.Config()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(restarted.Close) + pool.Close() + successor, _ := sessionExecution(t, restarted) + after, err := successor.PromoteEnvironmentInput(ctx, text(tenant), text(session), first.ID) + if err != nil || after.State != sessions.EnvironmentInputAdmitted || after.Receipts[0].Sequence != promoted.Receipts[0].Sequence { + t.Fatal("restart repeated promotion", after, err) + } + environmentInputHistory(t, restarted, session, 1, 2) +} + +// Reservations resolve only in their tenant's Session, take only messages, +// need an Environment and join an active Turn at once. +func TestEnvironmentInputReservationRejectsUnsupportedOrForeignState(t *testing.T) { + pool, store, service, operations := leasedInputs(t) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + ctx := t.Context() + pending := reserve(t, service, tenant, session, "pending") + stranger := pgID(uuid.New()) + for _, read := range []func() error{ + func() error { + _, err := store.GetEnvironmentInputReservation(ctx, text(stranger), text(session), pending.ID) + return err + }, + func() error { + _, err := operations.PromoteEnvironmentInput(ctx, text(stranger), text(session), pending.ID) + return err + }, + func() error { + _, err := cancelPending(ctx, pool, store, stranger, session, pending.ID) + return err + }, + func() error { + _, err := service.ReserveEnvironmentInput(ctx, text(stranger), text(session), "new", pending.Inputs) + return err + }, + } { + if err := read(); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal("foreign access", err) + } + } + other := pgID(uuid.New()) + exec(t, pool, `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash) VALUES ($1, $2, 'codex', 'other', 'hash')`, other, tenant) + for _, action := range []func(context.Context, pgtype.UUID, pgtype.UUID, string) (sessions.EnvironmentInputReservation, error){ + func(ctx context.Context, tenant, session pgtype.UUID, reservation string) (sessions.EnvironmentInputReservation, error) { + return store.GetEnvironmentInputReservation(ctx, text(tenant), text(session), reservation) + }, + func(ctx context.Context, tenant, session pgtype.UUID, reservation string) (sessions.EnvironmentInputReservation, error) { + return operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), reservation) + }, + func(ctx context.Context, tenant, session pgtype.UUID, reservation string) (sessions.EnvironmentInputReservation, error) { + return cancelPending(ctx, pool, store, tenant, session, reservation) + }, + func(ctx context.Context, tenant, session pgtype.UUID, reservation string) (sessions.EnvironmentInputReservation, error) { + return service.ExpireEnvironmentInput(ctx, text(tenant), text(session), reservation) + }, + } { + if _, err := action(ctx, tenant, other, pending.ID); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal("reservation crossed Session ownership", err) + } + if _, err := action(ctx, stranger, session, pending.ID); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal("reservation crossed tenant ownership", err) + } + } + retained, err := store.GetEnvironmentInputReservation(ctx, text(tenant), text(session), pending.ID) + if err != nil || !reflect.DeepEqual(retained, pending) { + t.Fatal("foreign operations changed reservation", retained, err) + } + for _, invalid := range [][]sessions.Input{nil, {cancelInput}, {{Kind: "tool_result", Payload: []byte(`{}`)}}} { + if _, err := service.ReserveEnvironmentInput(ctx, text(tenant), text(session), "invalid", invalid); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatal("unsupported reservation", err) + } + } + noneTenant, none := newSession(t, pool) + if _, err := service.ReserveEnvironmentInput(ctx, text(noneTenant), text(none), "none", pending.Inputs); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatal("none reservation", err) + } + activeTenant, active, _ := newEnvironment(t, pool, "self_hosted", "pending") + activeInput := submitMessage(t, service, activeTenant, active, "active") + steer, err := service.ReserveEnvironmentInput(ctx, text(activeTenant), text(active), "new", pending.Inputs) + if err != nil || steer.State != sessions.EnvironmentInputAdmitted || steer.ID != "" || !steer.Deadline.IsZero() || len(steer.Receipts) != len(pending.Inputs) || steer.Receipts[0].TurnID != activeInput.TurnID { + t.Fatal("active input did not retain the existing Turn", steer, err) + } +} + +// A cancelled or expired reservation keeps its outcome and identity: no retry +// or settlement restarts it, and none touches a newer reservation. +func TestEnvironmentInputTerminalReservationsCannotRestart(t *testing.T) { + pool, store, service, operations := leasedInputs(t) + for _, terminal := range []string{sessions.EnvironmentInputCancelled, sessions.EnvironmentInputExpired} { + t.Run(terminal, func(t *testing.T) { + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + ctx := t.Context() + pending := reserve(t, service, tenant, session, "pending") + early, err := service.ExpireEnvironmentInput(ctx, text(tenant), text(session), pending.ID) + if err != nil || early.State != sessions.EnvironmentInputPending || early.SettledAt != nil || !early.Deadline.Equal(pending.Deadline) { + t.Fatal("early expiry", early, err) + } + var settled sessions.EnvironmentInputReservation + if terminal == sessions.EnvironmentInputExpired { + passDeadline(t, pool, pending.ID) + settled, err = operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), pending.ID) + } else { + settled, err = cancelPending(ctx, pool, store, tenant, session, pending.ID) + } + if err != nil || settled.State != terminal || settled.SettledAt == nil || len(settled.Receipts) != 0 { + t.Fatal("terminal settlement", settled, err) + } + retry, err := service.ReserveEnvironmentInput(ctx, text(tenant), text(session), "pending", pending.Inputs) + if err != nil || retry.State != terminal || retry.ID != pending.ID || !retry.Deadline.Equal(settled.Deadline) || !retry.SettledAt.Equal(*settled.SettledAt) { + t.Fatal("terminal retry changed outcome", retry, err) + } + if _, err := service.SubmitInputs(ctx, text(tenant), text(session), "pending", pending.Inputs); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatal("terminal request reopened through direct path", err) + } + if _, err := service.SubmitInputs(ctx, text(tenant), text(session), "pending", []sessions.Input{messageInput("changed")}); !errors.Is(err, sessions.ErrIdempotencyConflict) { + t.Fatal("terminal identity changed", err) + } + later := reserve(t, service, tenant, session, "later") + // Cancellation takes whatever input is pending, so only the + // settlements that name the reservation run against it here. + for _, finish := range []func(context.Context, string, string, string) (sessions.EnvironmentInputReservation, error){ + operations.PromoteEnvironmentInput, service.ExpireEnvironmentInput, + } { + got, err := finish(ctx, text(tenant), text(session), pending.ID) + if err != nil || got.State != terminal { + t.Fatal("old settlement changed", got, err) + } + } + got, err := store.GetEnvironmentInputReservation(ctx, text(tenant), text(session), later.ID) + if err != nil || got.State != sessions.EnvironmentInputPending || !got.Deadline.Equal(later.Deadline) { + t.Fatal("old settlement touched successor", got, err) + } + environmentInputHistory(t, pool, session, 0, 0) + }) + } +} + +// A promotion that fails part way leaves no Turn, input, Item, claim or +// settlement, and its retry promotes the reservation whole. +func TestEnvironmentInputPromotionRollsBackHistoryAndSettlement(t *testing.T) { + pool, store, service, operations := leasedInputs(t) + for _, phase := range []string{"input", "settlement", "claim", "claim-event"} { + t.Run(phase, func(t *testing.T) { + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + ctx := t.Context() + pending := reserve(t, service, tenant, session, "pending") + name := "reservation_failure_" + strings.ReplaceAll(uuid.NewString(), "-", "") + table, expression := "turn_inputs", "session_id <> '"+text(session)+"'::uuid OR payload->>'text' <> 'second'" + switch phase { + case "settlement": + table, expression = "environment_input_reservations", "id <> '"+pending.ID+"'::uuid OR state <> 'admitted'" + case "claim": + table, expression = "turns", "session_id <> '"+text(session)+"'::uuid OR status <> 'in_progress'" + case "claim-event": + table, expression = "session_events", "session_id <> '"+text(session)+"'::uuid OR payload->'event'->>'type' <> 'agent.session.turn.in_progress'" + } + exec(t, pool, "ALTER TABLE "+table+" ADD CONSTRAINT "+name+" CHECK ("+expression+") NOT VALID") + t.Cleanup(func() { + _, _ = pool.Exec(context.Background(), "ALTER TABLE "+table+" DROP CONSTRAINT IF EXISTS "+name) + }) + if _, err := operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), pending.ID); err == nil { + t.Fatal("injected failure succeeded") + } + environmentInputHistory(t, pool, session, 0, 0) + got, err := store.GetEnvironmentInputReservation(ctx, text(tenant), text(session), pending.ID) + if err != nil || got.State != sessions.EnvironmentInputPending || got.SettledAt != nil || !got.Deadline.Equal(pending.Deadline) { + t.Fatal("partial settlement survived", got, err) + } + exec(t, pool, "ALTER TABLE "+table+" DROP CONSTRAINT "+name) + got, err = operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), pending.ID) + if err != nil || got.State != sessions.EnvironmentInputAdmitted { + t.Fatal(got, err) + } + environmentInputHistory(t, pool, session, 1, 2) + }) + } +} + +// Promotion and failure read the deadline after they take the Session lock, +// so waiting for the lock never extends the reservation's lifetime. +func TestEnvironmentInputDeadlineIsCheckedAfterSessionLock(t *testing.T) { + pool, store, service, operations := leasedInputs(t) + for _, action := range []string{"promote", "fail"} { + t.Run(action, func(t *testing.T) { + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + pending := reserve(t, service, tenant, session, "pending") + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + blocker, release := lockSession(ctx, t, pool, session) + type outcome struct { + value sessions.EnvironmentInputReservation + err error + } + done := make(chan outcome, 1) + go func() { + var got sessions.EnvironmentInputReservation + var err error + if action == "promote" { + got, err = operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), pending.ID) + } else if err = operations.FailEnvironmentInput(ctx, text(tenant), text(session), pending.ID, "runtime_preparation_failed"); err == nil { + got, err = store.GetEnvironmentInputReservation(ctx, text(tenant), text(session), pending.ID) + } + done <- outcome{got, err} + }() + awaitBlocked(ctx, t, pool, blocker) + // Transaction-start time is now older than the controlled deadline. + release("UPDATE environment_input_reservations SET deadline=clock_timestamp() WHERE id=$1", pending.ID) + result := <-done + if result.err != nil || result.value.State != sessions.EnvironmentInputExpired || result.value.SettledAt == nil { + t.Fatal("lock wait extended input lifetime", result) + } + environmentInputHistory(t, pool, session, 0, 0) + if state := reservationState(t, store, tenant, session, pending.ID); state != sessions.EnvironmentInputExpired { + t.Fatal("expiry was rolled back", state) + } + }) + } +} + +// A racing promotion and cancellation agree on one outcome. +func TestEnvironmentInputCancelAndPromotionShareOneOutcome(t *testing.T) { + pool, store, service, operations := leasedInputs(t) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + ctx := t.Context() + pending := reserve(t, service, tenant, session, "pending") + start := make(chan struct{}) + results := make(chan sessions.EnvironmentInputReservation, 2) + errs := make(chan error, 2) + for _, finish := range []func() (sessions.EnvironmentInputReservation, error){ + func() (sessions.EnvironmentInputReservation, error) { + return operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), pending.ID) + }, + func() (sessions.EnvironmentInputReservation, error) { + return cancelPending(ctx, pool, store, tenant, session, pending.ID) + }, + } { + go func() { + <-start + got, err := finish() + results <- got + errs <- err + }() + } + close(start) + first, second := <-results, <-results + for range 2 { + if err := <-errs; err != nil { + t.Fatal(err) + } + } + if first.State != second.State || (first.State != sessions.EnvironmentInputAdmitted && first.State != sessions.EnvironmentInputCancelled) { + t.Fatal("competing settlements diverged", first.State, second.State) + } + turns, inputs := 0, 0 + if first.State == sessions.EnvironmentInputAdmitted { + turns, inputs = 1, 2 + } + environmentInputHistory(t, pool, session, turns, inputs) +} + +// Each expiry pass settles at most 32 due reservations and creates no +// history. +func TestEnvironmentExpiryBoundsBatch(t *testing.T) { + pool, store, service, operations := leasedInputs(t) + ctx := t.Context() + type due struct { + tenant, session pgtype.UUID + reservation string + } + var reservations []due + for range 33 { + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + pending := reserve(t, service, tenant, session, "pending") + passDeadline(t, pool, pending.ID) + reservations = append(reservations, due{tenant, session, pending.ID}) + } + for _, want := range []int64{32, 1, 0} { + n, err := operations.ExpireEnvironmentInputs(ctx) + if err != nil || n != want { + t.Fatal("unbounded or incomplete batch", n, want, err) + } + } + for _, r := range reservations { + if state := reservationState(t, store, r.tenant, r.session, r.reservation); state != sessions.EnvironmentInputExpired { + t.Fatal(state) + } + environmentInputHistory(t, pool, r.session, 0, 0) + } +} + +// An execution owner that lost its lease expires nothing; its successor +// does. +func TestEnvironmentExpiryFencesLostExecutionOwner(t *testing.T) { + pool, store, service, old := leasedInputs(t) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + pending := reserve(t, service, tenant, session, "pending") + passDeadline(t, pool, pending.ID) + terminateLeaseOwner(t, pool) + successor, _ := sessionExecution(t, pool) + if n, err := old.ExpireEnvironmentInputs(t.Context()); err == nil || n != 0 { + t.Fatal("lost owner expired input", n, err) + } + if state := reservationState(t, store, tenant, session, pending.ID); state != sessions.EnvironmentInputPending { + t.Fatal("lost owner expired input", state) + } + if n, err := successor.ExpireEnvironmentInputs(t.Context()); err != nil || n != 1 { + t.Fatal("successor could not expire input", n, err) + } + if state := reservationState(t, store, tenant, session, pending.ID); state != sessions.EnvironmentInputExpired { + t.Fatal(state) + } + environmentInputHistory(t, pool, session, 0, 0) +} + +// Promotion and expiry of a reservation past its deadline, in either order or +// racing, end with one expired outcome and no history, and leave a newer +// reservation pending. +func TestEnvironmentInputPromotionSerializesWithExpiry(t *testing.T) { + pool, store, service, operations := leasedInputs(t) + promote := func(ctx context.Context, tenant, session pgtype.UUID, reservation string) error { + got, err := operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), reservation) + if err == nil && got.State != sessions.EnvironmentInputExpired { + return errors.New("promotion outran expiry: " + got.State) + } + return err + } + sweep := func(ctx context.Context, _, _ pgtype.UUID, _ string) error { + _, err := operations.ExpireEnvironmentInputs(ctx) + return err + } + expire := func(ctx context.Context, tenant, session pgtype.UUID, reservation string) error { + got, err := service.ExpireEnvironmentInput(ctx, text(tenant), text(session), reservation) + if err == nil && got.State != sessions.EnvironmentInputExpired { + return errors.New("expiry missed its deadline: " + got.State) + } + return err + } + type settle func(context.Context, pgtype.UUID, pgtype.UUID, string) error + for name, actions := range map[string][]settle{ + "sweep-then-promote": {sweep, promote}, + "promote-then-sweep": {promote, sweep}, + "racing-sweep": {sweep, promote}, + "racing-expire": {expire, promote}, + } { + t.Run(name, func(t *testing.T) { + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + ctx := t.Context() + pending := reserve(t, service, tenant, session, "pending") + passDeadline(t, pool, pending.ID) + if strings.HasPrefix(name, "racing") { + start := make(chan struct{}) + errs := make(chan error, len(actions)) + for _, action := range actions { + go func() { <-start; errs <- action(ctx, tenant, session, pending.ID) }() + } + close(start) + for range actions { + if err := <-errs; err != nil { + t.Fatal(err) + } + } + } else { + for _, action := range actions { + if err := action(ctx, tenant, session, pending.ID); err != nil { + t.Fatal(err) + } + } + } + if state := reservationState(t, store, tenant, session, pending.ID); state != sessions.EnvironmentInputExpired { + t.Fatal("invalid competing settlement", state) + } + environmentInputHistory(t, pool, session, 0, 0) + later := reserve(t, service, tenant, session, "later") + for _, action := range []settle{promote, expire, sweep} { + if err := action(ctx, tenant, session, pending.ID); err != nil { + t.Fatal("old reservation changed", err) + } + } + got, err := store.GetEnvironmentInputReservation(ctx, text(tenant), text(session), later.ID) + if err != nil || got.State != sessions.EnvironmentInputPending || !got.Deadline.Equal(later.Deadline) { + t.Fatal("old settlement affected successor", got, err) + } + }) + } +} + +// Concurrent promotions claim one Turn once: one fresh admission, one claim +// event after the admission's, and a retry after the Turn ends publishes +// nothing. +func TestEnvironmentInputConcurrentPromotionClaimsOnce(t *testing.T) { + pool, store, service, operations := leasedInputs(t) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + pending := reserve(t, service, tenant, session, "pending") + before, _ := journal(t, pool, session) + const count = 8 + results := make(chan sessions.EnvironmentInputReservation, count) + var group sync.WaitGroup + for range count { + group.Go(func() { + got, err := operations.PromoteEnvironmentInput(t.Context(), text(tenant), text(session), pending.ID) + if err != nil { + t.Error(err) + return + } + results <- got + }) + } + group.Wait() + close(results) + var turnID string + fresh, received := 0, 0 + for got := range results { + received++ + if got.State != sessions.EnvironmentInputAdmitted || len(got.Receipts) != 2 { + t.Fatal("promotion lost the original batch", got) + } + if turnID == "" { + turnID = got.Receipts[0].TurnID + } + if !got.Receipts[0].Replayed { + fresh++ + } + for _, receipt := range got.Receipts { + if receipt.TurnID != turnID || receipt.Replayed != got.Receipts[0].Replayed { + t.Fatal("promotion changed execution ownership", got.Receipts) + } + } + } + if received != count || fresh != 1 { + t.Fatal("promotion authorized multiple starts", received, fresh) + } + turn, err := store.GetTurn(t.Context(), text(tenant), text(session), turnID) + if err != nil || turn.Status != sessions.TurnInProgress || turn.StartedAt.IsZero() { + t.Fatal("promotion did not persist its execution claim", turn, err) + } + environmentInputHistory(t, pool, session, 1, 2) + _, changes := journal(t, pool, session) + changes = changes[len(before):] + if len(changes) < 2 { + t.Fatal("missing promotion events", kinds(changes)) + } + created, claimed := changes[0], changes[len(changes)-1] + if created.Event.Type != "agent.session.turn.created" || created.Turn == nil || created.Turn.Status != sessions.TurnQueued || claimed.Event.Type != "agent.session.turn.in_progress" || claimed.Turn == nil || claimed.Turn.ID != turnID || claimed.Turn.Status != sessions.TurnInProgress { + t.Fatal("claim reordered or replaced admission snapshots", kinds(changes)) + } + claims := 0 + for _, change := range changes { + if change.Event.Type == "agent.session.turn.in_progress" { + claims++ + } + } + if claims != 1 { + t.Fatal("retry published another claim", claims) + } + move(t, pool, tenant, session, turnID, sessions.TurnInProgress, sessions.TurnCompleted) + later := reserve(t, service, tenant, session, "later") + cursor, _ := journal(t, pool, session) + retry, err := operations.PromoteEnvironmentInput(t.Context(), text(tenant), text(session), pending.ID) + if err != nil || len(retry.Receipts) != 2 || !retry.Receipts[0].Replayed || retry.Receipts[0].TurnID != turnID { + t.Fatal("terminal retry reclaimed execution", retry, err) + } + if after, _ := journal(t, pool, session); !reflect.DeepEqual(after, cursor) { + t.Fatal("terminal retry published events", after, cursor) + } + retained, err := store.GetEnvironmentInputReservation(t.Context(), text(tenant), text(session), later.ID) + if err != nil || retained.State != sessions.EnvironmentInputPending || !retained.Deadline.Equal(later.Deadline) { + t.Fatal("old promotion affected new preparation", retained, err) + } +} + +// Promotion and preparation failure from a closed lease or a lost execution +// owner write nothing; the successor settles both. +func TestEnvironmentInputSettlementUsesCurrentExecutionOwner(t *testing.T) { + pool := pgtest.OpenIsolated(t, nil) + store, service := stagingService(t, pool) + promoteTenant, promoteSession, _ := newEnvironment(t, pool, "self_hosted", "pending") + failTenant, failSession, _ := newEnvironment(t, pool, "self_hosted", "pending") + promoted := reserve(t, service, promoteTenant, promoteSession, "pending") + failed := reserve(t, service, failTenant, failSession, "pending") + settle := func(operations *sessions.ExecutionOperations) (error, error) { + _, promoteErr := operations.PromoteEnvironmentInput(t.Context(), text(promoteTenant), text(promoteSession), promoted.ID) + return promoteErr, operations.FailEnvironmentInput(t.Context(), text(failTenant), text(failSession), failed.ID, "runtime_preparation_failed") + } + unsettled := func() { + t.Helper() + environmentInputHistory(t, pool, promoteSession, 0, 0) + environmentInputHistory(t, pool, failSession, 0, 0) + if reservationState(t, store, promoteTenant, promoteSession, promoted.ID) != sessions.EnvironmentInputPending || reservationState(t, store, failTenant, failSession, failed.ID) != sessions.EnvironmentInputPending { + t.Fatal("fenced owner settled input") + } + } + closed, lease := sessionExecution(t, pool) + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, pool) + if err := lease.Close(t.Context()); err != nil { + t.Fatal(err) + } + awaitRelease() + if promoteErr, failErr := settle(closed); !errors.Is(promoteErr, pgunit.ErrLeaseClosed) || !errors.Is(failErr, pgunit.ErrLeaseClosed) { + t.Fatal("closed lease settled input", promoteErr, failErr) + } + unsettled() + stale, _ := sessionExecution(t, pool) + terminateLeaseOwner(t, pool) + successor, _ := sessionExecution(t, pool) + if promoteErr, failErr := settle(stale); promoteErr == nil || failErr == nil { + t.Fatal("stale owner settled input", promoteErr, failErr) + } + unsettled() + if promoteErr, failErr := settle(successor); promoteErr != nil || failErr != nil { + t.Fatal("successor could not settle", promoteErr, failErr) + } + environmentInputHistory(t, pool, promoteSession, 1, 2) + environmentInputHistory(t, pool, failSession, 0, 0) + if reservationState(t, store, promoteTenant, promoteSession, promoted.ID) != sessions.EnvironmentInputAdmitted || reservationState(t, store, failTenant, failSession, failed.ID) != sessions.EnvironmentInputFailed { + t.Fatal("successor settlement") + } +} + +// Input reserved while the active Turn completes serializes with the +// completion: input first joins the Turn and blocks completion until applied; +// completion first leaves the input pending for its own Turn. +func TestEnvironmentActiveInputSerializesWithCompletion(t *testing.T) { + pool, _, service, operations := leasedInputs(t) + for _, completionFirst := range []bool{false, true} { + name := "input-first" + if completionFirst { + name = "completion-first" + } + t.Run(name, func(t *testing.T) { + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + original := submitMessage(t, service, tenant, session, "original") + move(t, pool, tenant, session, original.TurnID, sessions.TurnQueued, sessions.TurnInProgress) + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + blocker, release := lockSession(ctx, t, pool, session) + type admission struct { + value sessions.EnvironmentInputReservation + err error + } + admitted := make(chan admission, 1) + completed := make(chan error, 1) + input := func() { + value, err := service.ReserveEnvironmentInput(ctx, text(tenant), text(session), "racing-input", reservedBatch) + admitted <- admission{value, err} + } + complete := func() { + _, err := operations.CompleteExecution(ctx, text(tenant), text(session), original.TurnID, sessions.TurnCompleted, nil, "", original.Sequence) + completed <- err + } + first, second := input, complete + if completionFirst { + first, second = complete, input + } + go first() + firstPID := awaitBlocked(ctx, t, pool, blocker) + go second() + awaitBlocked(ctx, t, pool, firstPID) + release("") + got, completionErr := <-admitted, <-completed + if got.err != nil { + t.Fatal(got.err) + } + if completionFirst { + if completionErr != nil || got.value.State != sessions.EnvironmentInputPending || got.value.ID == "" || len(got.value.Receipts) != 0 || got.value.Deadline.Sub(got.value.CreatedAt) != 5*time.Minute { + t.Fatal("completion winner did not leave new input waiting for preparation", completionErr, got.value) + } + environmentInputHistory(t, pool, session, 1, 1) + prepared, err := operations.PromoteEnvironmentInput(ctx, text(tenant), text(session), got.value.ID) + if err != nil || len(prepared.Receipts) != 2 || prepared.Receipts[0].TurnID == original.TurnID { + t.Fatal("prepared successor reused terminal work", err) + } + retry, err := service.ReserveEnvironmentInput(ctx, text(tenant), text(session), "racing-input", reservedBatch) + if err != nil || retry.ID != got.value.ID || !retry.Deadline.Equal(got.value.Deadline) || len(retry.Receipts) != 2 || !retry.Receipts[0].Replayed || retry.Receipts[0].TurnID != prepared.Receipts[0].TurnID { + t.Fatal("active retry replaced its original reservation", retry, err) + } + if _, err := operations.CompleteExecution(ctx, text(tenant), text(session), prepared.Receipts[0].TurnID, sessions.TurnCompleted, nil, "", prepared.Receipts[1].Sequence); err != nil { + t.Fatal(err) + } + } else { + if !errors.Is(completionErr, sessions.ErrUnappliedInputs) || got.value.State != sessions.EnvironmentInputAdmitted || got.value.ID != "" || !got.value.Deadline.IsZero() || len(got.value.Receipts) != 2 || got.value.Receipts[0].TurnID != original.TurnID || got.value.Receipts[0].Replayed { + t.Fatal("admitted input escaped the original Turn or application fence", completionErr, got.value) + } + environmentInputHistory(t, pool, session, 1, 3) + if _, err := operations.CompleteExecution(ctx, text(tenant), text(session), original.TurnID, sessions.TurnCompleted, nil, "", got.value.Receipts[1].Sequence); err != nil { + t.Fatal("completion after controlled application failed", err) + } + } + }) + } +} + +// A late preparation failure leaves a cancelled reservation and a newer one +// as they are, and takes only classified codes. +func TestPreparationFailurePreservesCancelledAndNewerInput(t *testing.T) { + pool, store, service, operations := leasedInputs(t) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "pending") + first := reserve(t, service, tenant, session, "cancelled") + if _, err := cancelPending(t.Context(), pool, store, tenant, session, first.ID); err != nil { + t.Fatal(err) + } + next := reserve(t, service, tenant, session, "new") + if err := operations.FailEnvironmentInput(t.Context(), text(tenant), text(session), first.ID, "runtime_preparation_failed"); err != nil { + t.Fatal(err) + } + for id, state := range map[string]string{first.ID: sessions.EnvironmentInputCancelled, next.ID: sessions.EnvironmentInputPending} { + if current := reservationState(t, store, tenant, session, id); current != state { + t.Fatal("late failure changed another outcome", id, current) + } + } + if err := operations.FailEnvironmentInput(t.Context(), text(tenant), text(session), next.ID, "secret-canary"); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatal("unclassified diagnostic accepted", err) + } +} diff --git a/services/core/internal/persistence/postgres/sessionpg/inputs.go b/services/core/internal/persistence/postgres/sessionpg/inputs.go new file mode 100644 index 000000000..37e764773 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/inputs.go @@ -0,0 +1,247 @@ +package sessionpg + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var ( + _ sessions.InputTx = (*SessionTx)(nil) + _ sessions.InputReader = (*Store)(nil) +) + +func (s *Store) WithInputs(ctx context.Context, tenant, session string, apply func(context.Context, sessions.InputTx) error) error { + return withInputs(ctx, s.units, tenant, session, apply) +} + +// withInputs runs apply on runner in the Session transaction of the tenant's +// visible Session. A malformed tenant is sessions.ErrInvalidInput; a malformed +// Session ID resolves as a missing Session. +func withInputs(ctx context.Context, runner pgunit.Transactor, tenantID, sessionID string, apply func(context.Context, sessions.InputTx) error) error { + tenant, err := parseID(tenantID) + if err != nil { + return err + } + session := pgunit.PathID(sessionID) + return WithSession(ctx, runner, tenant, session, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + if err := locked.Public(); err != nil { + return err + } + return apply(ctx, BindSession(q, tenant, session)) + }) +} + +func (t *SessionTx) CreateTurn(ctx context.Context) (sessions.Turn, error) { + row, err := t.q.CreateTurn(ctx, sqlc.CreateTurnParams{ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, SessionID: t.session}) + if err != nil { + return sessions.Turn{}, err + } + return TurnFromRow(row), nil +} + +func (t *SessionTx) CreateTurnInput(ctx context.Context, turn, key string, position int32, input sessions.Input) (int64, error) { + var id pgtype.UUID + if turn != "" { + var err error + if id, err = parseID(turn); err != nil { + return 0, err + } + } + sequence, err := t.q.CreateTurnInput(ctx, sqlc.CreateTurnInputParams{ + SessionID: t.session, TurnID: id, IdempotencyKey: key, Kind: input.Kind, Payload: input.Payload, BatchPosition: position, + }) + return sequence, storable(err) +} + +func (t *SessionTx) LoadInputBatch(ctx context.Context, key string, batch json.RawMessage) ([]sessions.InputReceipt, bool, error) { + rows, err := t.q.FindInputBatch(ctx, sqlc.FindInputBatchParams{SessionID: t.session, IdempotencyKey: key, Batch: batch}) + if err != nil { + return nil, false, storable(err) + } + receipts := make([]sessions.InputReceipt, 0, len(rows)) + matches := true + for _, row := range rows { + matches = matches && row.Matches + receipts = append(receipts, sessions.InputReceipt{Sequence: row.Sequence, TurnID: uuidString(row.TurnID), Replayed: true}) + } + return receipts, matches, nil +} + +func (t *SessionTx) LoadInputGate(ctx context.Context, key string, batch json.RawMessage) (sessions.InputGate, error) { + gate, err := t.q.CheckEnvironmentInputGate(ctx, sqlc.CheckEnvironmentInputGateParams{SessionID: t.session, IdempotencyKey: key, Batch: batch}) + return sessions.InputGate{Matches: gate.Matches, Blocked: gate.Blocked}, storable(err) +} + +func (t *SessionTx) FindInputReservation(ctx context.Context, key string, batch json.RawMessage) (*sessions.EnvironmentInputReservation, bool, error) { + row, err := t.q.FindEnvironmentInputReservation(ctx, sqlc.FindEnvironmentInputReservationParams{SessionID: t.session, IdempotencyKey: key, Batch: batch}) + if errors.Is(err, pgx.ErrNoRows) { + return nil, false, nil + } + if err != nil { + return nil, false, storable(err) + } + if !row.Matches { + return nil, true, nil + } + reservation, err := t.reservationFromRow(ctx, row.EnvironmentInputReservation) + if err != nil { + return nil, false, err + } + return &reservation, true, nil +} + +func (t *SessionTx) LoadInputReservation(ctx context.Context, reservation string) (sessions.EnvironmentInputReservation, error) { + id, err := parseID(reservation) + if err != nil { + return sessions.EnvironmentInputReservation{}, err + } + row, err := t.q.GetEnvironmentInputReservation(ctx, sqlc.GetEnvironmentInputReservationParams{SessionID: t.session, ID: id}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.EnvironmentInputReservation{}, sessions.ErrNotFound + } + if err != nil { + return sessions.EnvironmentInputReservation{}, err + } + return t.reservationFromRow(ctx, row) +} + +// reservationFromRow maps a stored Environment input reservation; an +// admitted one carries the receipts of its batch, which it must have. +func (t *SessionTx) reservationFromRow(ctx context.Context, row sqlc.EnvironmentInputReservation) (sessions.EnvironmentInputReservation, error) { + result := sessions.EnvironmentInputReservation{ + ID: uuid.UUID(row.ID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), Key: row.IdempotencyKey, + State: row.State, IsInitial: row.IsInitial, CreatedAt: row.CreatedAt.Time, Deadline: row.Deadline.Time, + } + if row.SettledAt.Valid { + result.SettledAt = &row.SettledAt.Time + } + if err := json.Unmarshal(row.Batch, &result.Inputs); err != nil { + return sessions.EnvironmentInputReservation{}, fmt.Errorf("decode Environment input: %w", err) + } + if row.State != sessions.EnvironmentInputAdmitted { + return result, nil + } + receipts, matches, err := t.LoadInputBatch(ctx, row.IdempotencyKey, row.Batch) + if err != nil { + return sessions.EnvironmentInputReservation{}, err + } + if len(receipts) == 0 || !matches { + return sessions.EnvironmentInputReservation{}, errors.New("admitted Environment input has no receipts of its batch") + } + result.Receipts = receipts + return result, nil +} + +func (t *SessionTx) CreateInputReservation(ctx context.Context, key string, batch json.RawMessage) (sessions.EnvironmentInputReservation, error) { + row, err := t.q.CreateEnvironmentInputReservation(ctx, sqlc.CreateEnvironmentInputReservationParams{ + ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, SessionID: t.session, IdempotencyKey: key, Batch: batch, + }) + if err != nil { + return sessions.EnvironmentInputReservation{}, storable(err) + } + return t.reservationFromRow(ctx, row) +} + +func (t *SessionTx) ExpireInputReservation(ctx context.Context, reservation string) error { + id, err := parseID(reservation) + if err != nil { + return err + } + return t.q.ExpireEnvironmentInputReservation(ctx, sqlc.ExpireEnvironmentInputReservationParams{SessionID: t.session, ID: id}) +} + +func (t *SessionTx) AdmitInputReservation(ctx context.Context, reservation string) (time.Time, error) { + id, err := parseID(reservation) + if err != nil { + return time.Time{}, err + } + row, err := t.q.SettleEnvironmentInputReservation(ctx, sqlc.SettleEnvironmentInputReservationParams{SessionID: t.session, ID: id, State: sessions.EnvironmentInputAdmitted}) + if err != nil { + return time.Time{}, err + } + return row.SettledAt.Time, nil +} + +func (t *SessionTx) FailInputReservation(ctx context.Context, reservation, code string) error { + id, err := parseID(reservation) + if err != nil { + return err + } + _, err = t.q.FailEnvironmentInput(ctx, sqlc.FailEnvironmentInputParams{SessionID: t.session, ID: id, FailureCode: pgtype.Text{String: code, Valid: true}}) + return err +} + +func (t *SessionTx) RecordInputAudit(ctx context.Context) error { + return auditpg.RecordWriteAudit(ctx, t.q, optionalID(t.tenant), "send_events", "session", optionalID(t.session), "") +} + +func (s *Store) ListTurnInputs(ctx context.Context, tenant, session, turn string, after int64, limit int) ([]sessions.TurnInput, error) { + params, err := TurnLookup(tenant, session, turn) + if err != nil { + return nil, err + } + if after < 0 || limit < 1 || limit > 100 { + return nil, fmt.Errorf("%w: nonnegative cursor and page size 1..100 required", sessions.ErrInvalidInput) + } + if _, err := s.GetTurn(ctx, tenant, session, turn); err != nil { + return nil, err + } + rows, err := s.units.Queries().ListTurnInputs(ctx, sqlc.ListTurnInputsParams{ + TenantID: params.TenantID, SessionID: params.SessionID, TurnID: params.ID, Sequence: after, Limit: int32(limit), + }) + if err != nil { + return nil, fmt.Errorf("list turn inputs: %w", err) + } + inputs := make([]sessions.TurnInput, 0, len(rows)) + for _, row := range rows { + inputs = append(inputs, sessions.TurnInput{Sequence: row.Sequence, Kind: row.Kind, Payload: row.Payload, CreatedAt: row.CreatedAt.Time}) + } + return inputs, nil +} + +func (s *Store) GetEnvironmentInputReservation(ctx context.Context, tenantID, sessionID, reservationID string) (sessions.EnvironmentInputReservation, error) { + tenant, err := parseID(tenantID) + if err != nil { + return sessions.EnvironmentInputReservation{}, err + } + session := pgunit.PathID(sessionID) + var reservation sessions.EnvironmentInputReservation + err = s.units.Snapshot(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + if err := visibleSession(ctx, q, tenant, session); err != nil { + return err + } + var err error + reservation, err = BindSession(q, tenant, session).LoadInputReservation(ctx, reservationID) + return err + }) + return reservation, err +} + +func (s *Store) ListEnvironmentInputWork(ctx context.Context, after string, connectedDevices []string) ([]sessions.EnvironmentInputWork, error) { + id, devices, err := executionWorkCursor(after, connectedDevices) + if err != nil { + return nil, err + } + rows, err := s.units.Queries().ListEnvironmentInputWork(ctx, sqlc.ListEnvironmentInputWorkParams{AfterID: id, ConnectedDevices: devices}) + if err != nil { + return nil, err + } + work := make([]sessions.EnvironmentInputWork, 0, len(rows)) + for _, row := range rows { + work = append(work, sessions.EnvironmentInputWork{TenantID: uuid.UUID(row.TenantID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), ReservationID: uuid.UUID(row.ID.Bytes).String()}) + } + return work, nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/inputs_test.go b/services/core/internal/persistence/postgres/sessionpg/inputs_test.go new file mode 100644 index 000000000..e2f9e1b2f --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/inputs_test.go @@ -0,0 +1,521 @@ +package sessionpg + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "reflect" + "strings" + "sync" + "testing" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var messagePayload = json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"hello"}]}]}`) + +var cancelInput = sessions.Input{Kind: "cancel", Payload: json.RawMessage(`{}`)} + +func messageInput(text string) sessions.Input { + payload, _ := json.Marshal(map[string]string{"text": text}) + return sessions.Input{Kind: "message", Payload: payload} +} + +// newSession stores a fresh tenant's idle Session without an Environment and +// returns their IDs. +func newSession(t *testing.T, pool *pgxpool.Pool) (pgtype.UUID, pgtype.UUID) { + t.Helper() + tenant, session := uuid.New(), uuid.New() + exec(t, pool, `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash) VALUES ($1, $2, 'codex', 'key', 'hash')`, session, tenant) + return pgID(tenant), pgID(session) +} + +// submit admits one input through the Session service. +func submit(ctx context.Context, service *sessions.Service, tenant, session pgtype.UUID, key string, input sessions.Input) (sessions.InputReceipt, error) { + receipts, err := service.SubmitInputs(ctx, text(tenant), text(session), key, []sessions.Input{input}) + if err != nil { + return sessions.InputReceipt{}, err + } + return receipts[0], nil +} + +func submitMessage(t *testing.T, service *sessions.Service, tenant, session pgtype.UUID, key string) sessions.InputReceipt { + t.Helper() + receipt, err := submit(t.Context(), service, tenant, session, key, sessions.Input{Kind: "message", Payload: messagePayload}) + if err != nil { + t.Fatal(err) + } + return receipt +} + +func move(t *testing.T, pool *pgxpool.Pool, tenant, session pgtype.UUID, turn, from, to string) sessions.Turn { + t.Helper() + moved, err := transition(t.Context(), pool, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: from, Status: to}) + if err != nil { + t.Fatal(err) + } + return moved +} + +// Concurrent messages from two pools all join one Turn, and concurrent retries +// of one key commit one input. +func TestConcurrentInputsUseOneTurnAndOneRetryReceipt(t *testing.T) { + pool := pgtest.Open(t) + store, service := stagingService(t, pool) + _, other := stagingService(t, pgtest.Open(t)) + tenant, session := newSession(t, pool) + const count = 8 + for _, repeatedKey := range []bool{true, false} { + t.Run(fmt.Sprintf("repeated-key-%v", repeatedKey), func(t *testing.T) { + var wg sync.WaitGroup + receipts := make(chan sessions.InputReceipt, count) + errs := make(chan error, count) + for i := range count { + wg.Go(func() { + key := "same" + if !repeatedKey { + key = fmt.Sprintf("distinct-%d", i) + } + admission := service + if i%2 == 0 { + admission = other + } + r, err := submit(t.Context(), admission, tenant, session, key, sessions.Input{Kind: "message", Payload: messagePayload}) + receipts <- r + errs <- err + }) + } + wg.Wait() + close(receipts) + close(errs) + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + turns, sequences := map[string]bool{}, map[int64]bool{} + newReceipts := 0 + for r := range receipts { + turns[r.TurnID], sequences[r.Sequence] = true, true + if !r.Replayed { + newReceipts++ + } + } + want := count + if repeatedKey { + want = 1 + } + if len(turns) != 1 || turns[""] || len(sequences) != want || newReceipts != want { + t.Fatalf("turns=%v sequences=%v new=%d", turns, sequences, newReceipts) + } + }) + } + first := submitMessage(t, service, tenant, session, "same") + inputs, err := store.ListTurnInputs(t.Context(), text(tenant), text(session), first.TurnID, 0, 100) + if err != nil || len(inputs) != count+1 { + t.Fatalf("lost or duplicated inputs: %d, %v", len(inputs), err) + } +} + +// An active Turn takes steering input, retries replay their receipts across a +// restart, an idle Session starts a new Turn, and admission changes no stored +// Session field besides its event sequence. +func TestTurnInputRetriesAndRestart(t *testing.T) { + pool := pgtest.Open(t) + _, service := stagingService(t, pool) + tenant, session := newSession(t, pool) + ctx := t.Context() + stored := func(pool *pgxpool.Pool) (row string) { + if err := pool.QueryRow(ctx, `SELECT (to_jsonb(s) - 'event_sequence')::text FROM sessions s WHERE id = $1`, session).Scan(&row); err != nil { + t.Fatal(err) + } + return row + } + created := stored(pool) + first := submitMessage(t, service, tenant, session, "first") + move(t, pool, tenant, session, first.TurnID, sessions.TurnQueued, sessions.TurnInProgress) + steer := submitMessage(t, service, tenant, session, "steer") + if steer.TurnID != first.TurnID || steer.Sequence <= first.Sequence { + t.Fatalf("active message did not steer: %+v", steer) + } + reordered := json.RawMessage(`{ "input": [{"content":[{"text":"hello","type":"input_text"}],"role":"user"}] }`) + retry, err := submit(ctx, service, tenant, session, "first", sessions.Input{Kind: "message", Payload: reordered}) + if err != nil || !retry.Replayed || retry.Sequence != first.Sequence || retry.TurnID != first.TurnID { + t.Fatalf("equivalent retry = %+v, %v", retry, err) + } + if _, err := submit(ctx, service, tenant, session, "first", messageInput("changed")); !errors.Is(err, sessions.ErrIdempotencyConflict) { + t.Fatalf("changed payload accepted: %v", err) + } + if _, err := submit(ctx, service, tenant, session, "first", cancelInput); !errors.Is(err, sessions.ErrIdempotencyConflict) { + t.Fatalf("changed input kind accepted: %v", err) + } + completed := move(t, pool, tenant, session, first.TurnID, sessions.TurnInProgress, sessions.TurnCompleted) + next := submitMessage(t, service, tenant, session, "next") + if next.TurnID == first.TurnID { + t.Fatal("idle message did not start a new Turn") + } + pool.Close() + reopened := pgtest.Open(t) + recovered, restarted := stagingService(t, reopened) + retry = submitMessage(t, restarted, tenant, session, "first") + if !retry.Replayed || retry.TurnID != first.TurnID || retry.Sequence != first.Sequence { + t.Fatalf("restart retry changed target: %+v", retry) + } + got, err := recovered.GetTurn(ctx, text(tenant), text(session), first.TurnID) + if err != nil || !reflect.DeepEqual(got, completed) { + t.Fatalf("restart turn: %+v, %v", got, err) + } + var all []sessions.TurnInput + var cursor int64 + for { + page, err := recovered.ListTurnInputs(ctx, text(tenant), text(session), first.TurnID, cursor, 1) + if err != nil { + t.Fatal(err) + } + if len(page) == 0 { + break + } + if page[0].Sequence <= cursor || page[0].CreatedAt.IsZero() { + t.Fatalf("invalid ordered input: %+v", page) + } + cursor = page[0].Sequence + all = append(all, page...) + } + if len(all) != 2 || all[0].Sequence != first.Sequence || all[1].Sequence != steer.Sequence { + t.Fatalf("recovered inputs = %+v", all) + } + snapshot, err := recovered.GetSession(ctx, text(tenant), text(session)) + if err != nil || snapshot.LastTurn == nil || snapshot.LastTurn.ID != next.TurnID || snapshot.LastTurn.Status != sessions.TurnQueued { + t.Fatal("latest Session activity did not survive restart", err) + } + if stored(reopened) != created { + t.Fatal("input admission mutated the stored Session") + } +} + +// Input admission and Turn reads and writes resolve only the tenant's own +// Session, and a Turn only within its Session. +func TestTurnOperationsAreTenantAndSessionScoped(t *testing.T) { + pool := pgtest.Open(t) + store, service := stagingService(t, pool) + tenant, session := newSession(t, pool) + otherTenant, other := newSession(t, pool) + ctx := t.Context() + first := submitMessage(t, service, tenant, session, "input") + for _, scope := range []struct{ tenant, session pgtype.UUID }{{otherTenant, session}, {tenant, other}, {tenant, pgID(uuid.New())}} { + for name, call := range map[string]func() error{ + "submit": func() error { + _, err := submit(ctx, service, scope.tenant, scope.session, "input", sessions.Input{Kind: "message", Payload: messagePayload}) + return err + }, + "cancel": func() error { + _, err := submit(ctx, service, scope.tenant, scope.session, "cancel", cancelInput) + return err + }, + "read": func() error { + _, err := store.GetTurn(ctx, text(scope.tenant), text(scope.session), first.TurnID) + return err + }, + "inputs": func() error { + _, err := store.ListTurnInputs(ctx, text(scope.tenant), text(scope.session), first.TurnID, 0, 10) + return err + }, + "transition": func() error { + _, err := transition(ctx, pool, scope.tenant, scope.session, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed}) + return err + }, + } { + if err := call(); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("%s escaped scope: %v", name, err) + } + } + } + // Turn IDs cannot be used with another valid Session in the same tenant either. + second := pgID(uuid.New()) + exec(t, pool, `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash) VALUES ($1, $2, 'codex', 'second', 'hash')`, second, tenant) + if _, err := store.GetTurn(ctx, text(tenant), text(second), first.TurnID); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("cross-session turn read: %v", err) + } + if _, err := transition(ctx, pool, tenant, second, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed}); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("cross-session turn write: %v", err) + } + otherInput := submitMessage(t, service, otherTenant, other, "input") + if otherInput.TurnID == first.TurnID { + t.Fatal("retry identity leaked across Sessions") + } +} + +// Batches from two pools keep their inputs adjacent and in order, and one +// key's concurrent retries commit one batch. +func TestInputBatchesAreOrderedAndIdempotentAcrossConnections(t *testing.T) { + pool := pgtest.Open(t) + store, service := stagingService(t, pool) + _, other := stagingService(t, pgtest.Open(t)) + tenant, session := newSession(t, pool) + ctx := t.Context() + const count = 8 + batch := []sessions.Input{messageInput("first"), messageInput("second")} + for _, repeated := range []bool{true, false} { + var wg sync.WaitGroup + receipts := make(chan []sessions.InputReceipt, count) + errs := make(chan error, count) + for i := range count { + wg.Go(func() { + admission := service + if i%2 == 0 { + admission = other + } + key := "same-batch" + if !repeated { + key = fmt.Sprintf("batch-%d", i) + } + got, err := admission.SubmitInputs(ctx, text(tenant), text(session), key, batch) + receipts <- got + errs <- err + }) + } + wg.Wait() + close(receipts) + close(errs) + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + newBatches := 0 + for got := range receipts { + if len(got) != 2 || got[0].TurnID == "" || got[0].TurnID != got[1].TurnID || got[0].Sequence >= got[1].Sequence || got[0].Replayed != got[1].Replayed { + t.Fatalf("invalid batch receipt: %+v", got) + } + if !got[0].Replayed { + newBatches++ + } + } + want := count + if repeated { + want = 1 + } + if newBatches != want { + t.Fatalf("accepted %d batches, want %d", newBatches, want) + } + } + got, err := service.SubmitInputs(ctx, text(tenant), text(session), "same-batch", batch) + if err != nil { + t.Fatal(err) + } + inputs, err := store.ListTurnInputs(ctx, text(tenant), text(session), got[0].TurnID, 0, 100) + if err != nil || len(inputs) != 2*(count+1) { + t.Fatalf("inputs=%d err=%v", len(inputs), err) + } + for i, input := range inputs { + var payload map[string]string + if err := json.Unmarshal(input.Payload, &payload); err != nil { + t.Fatal(err) + } + want := "first" + if i%2 == 1 { + want = "second" + } + if payload["text"] != want { + t.Fatalf("batch interleaved at %d: %v", i, payload) + } + } +} + +// A batch retry compares the whole normalized request and replays the +// cancellation targets it committed, after a restart too. +func TestBatchRetriesCompareTheWholeRequestAndRetainTargets(t *testing.T) { + pool := pgtest.Open(t) + store, service := stagingService(t, pool) + tenant, session := newSession(t, pool) + ctx := t.Context() + batch := []sessions.Input{cancelInput, messageInput("one"), cancelInput, messageInput("two")} + first, err := service.SubmitInputs(ctx, text(tenant), text(session), "mixed", batch) + if err != nil { + t.Fatal(err) + } + if first[0].TurnID != "" || first[1].TurnID == "" || first[1].TurnID != first[2].TurnID || first[1].TurnID == first[3].TurnID { + t.Fatalf("cancellation targets: %+v", first) + } + cancelled, err := store.GetTurn(ctx, text(tenant), text(session), first[1].TurnID) + if err != nil || cancelled.Status != sessions.TurnCancelled { + t.Fatalf("cancelled turn=%+v err=%v", cancelled, err) + } + move(t, pool, tenant, session, first[3].TurnID, sessions.TurnQueued, sessions.TurnInProgress) + move(t, pool, tenant, session, first[3].TurnID, sessions.TurnInProgress, sessions.TurnCompleted) + next := submitMessage(t, service, tenant, session, "next") + for _, changed := range [][]sessions.Input{batch[:3], append(append([]sessions.Input{}, batch...), cancelInput), {batch[0], batch[3], batch[2], batch[1]}} { + if _, err := service.SubmitInputs(ctx, text(tenant), text(session), "mixed", changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { + t.Fatalf("changed batch accepted: %v", err) + } + } + batch[1].Payload = json.RawMessage(`{ "text" : "one" }`) + pool.Close() + restartedStore, restarted := stagingService(t, pgtest.Open(t)) + retry, err := restarted.SubmitInputs(ctx, text(tenant), text(session), "mixed", batch) + for i := range first { + first[i].Replayed = true + } + if err != nil || !reflect.DeepEqual(retry, first) { + t.Fatalf("restart changed receipts: %+v, %v", retry, err) + } + current, err := restartedStore.GetTurn(ctx, text(tenant), text(session), next.TurnID) + if err != nil || current.Status != sessions.TurnQueued || !current.CancelRequestedAt.IsZero() { + t.Fatalf("retry cancelled later work: %+v, %v", current, err) + } + if _, err := restarted.SubmitInputs(ctx, uuid.NewString(), text(session), "mixed", batch); !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("batch retry escaped tenant: %v", err) + } +} + +// A batch that fails part way rolls back its earlier cancellation and inputs, +// and its retry then commits it whole. +func TestFailedBatchRollsBackEarlierCancellationAndInputs(t *testing.T) { + pool := pgtest.Open(t) + store, service := stagingService(t, pool) + tenant, session := newSession(t, pool) + ctx := t.Context() + initial := submitMessage(t, service, tenant, session, "initial") + move(t, pool, tenant, session, initial.TurnID, sessions.TurnQueued, sessions.TurnInProgress) + key := uuid.NewString() + constraint := "batch_failure_" + strings.ReplaceAll(key, "-", "") + // Inject a storage error on the second insert, after cancellation has run. + exec(t, pool, "ALTER TABLE turn_inputs ADD CONSTRAINT "+constraint+" CHECK (idempotency_key <> '"+key+"' OR batch_position = 0)") + t.Cleanup(func() { + _, _ = pool.Exec(context.Background(), "ALTER TABLE turn_inputs DROP CONSTRAINT IF EXISTS "+constraint) + }) + batch := []sessions.Input{cancelInput, messageInput("after cancel")} + if got, err := service.SubmitInputs(ctx, text(tenant), text(session), key, batch); err == nil || got != nil { + t.Fatalf("partial batch succeeded: %+v %v", got, err) + } + turn, err := store.GetTurn(ctx, text(tenant), text(session), initial.TurnID) + if err != nil || turn.Status != sessions.TurnInProgress || !turn.CancelRequestedAt.IsZero() { + t.Fatalf("cancellation escaped rollback: %+v %v", turn, err) + } + inputs, err := store.ListTurnInputs(ctx, text(tenant), text(session), initial.TurnID, 0, 100) + if err != nil || len(inputs) != 1 { + t.Fatalf("partial inputs survived: %v %v", inputs, err) + } + exec(t, pool, "ALTER TABLE turn_inputs DROP CONSTRAINT "+constraint) + got, err := service.SubmitInputs(ctx, text(tenant), text(session), key, batch) + if err != nil || len(got) != 2 || got[0].Replayed || got[1].Replayed { + t.Fatalf("retry after rollback: %+v %v", got, err) + } +} + +// A cancel request targets the Turn active when it was admitted, and its retry +// never retargets a later Turn; cancelling queued work stops it at once. +func TestCancellationStaysBoundToItsOriginalTurn(t *testing.T) { + pool := pgtest.Open(t) + store, service := stagingService(t, pool) + tenant, session := newSession(t, pool) + ctx := t.Context() + idle, err := submit(ctx, service, tenant, session, "idle-cancel", cancelInput) + if err != nil || idle.TurnID != "" || idle.Sequence == 0 { + t.Fatalf("idle cancellation: %+v, %v", idle, err) + } + first := submitMessage(t, service, tenant, session, "first") + started := move(t, pool, tenant, session, first.TurnID, sessions.TurnQueued, sessions.TurnInProgress) + if started.StartedAt.IsZero() || !started.CompletedAt.IsZero() { + t.Fatalf("started timestamps: %+v", started) + } + cancel, err := submit(ctx, service, tenant, session, "cancel", cancelInput) + if err != nil || cancel.TurnID != first.TurnID { + t.Fatalf("cancellation target: %+v, %v", cancel, err) + } + pending, err := store.GetTurn(ctx, text(tenant), text(session), first.TurnID) + if err != nil || pending.Status != sessions.TurnInProgress || pending.CancelRequestedAt.IsZero() || !pending.CompletedAt.IsZero() { + t.Fatalf("running cancellation falsely completed: %+v, %v", pending, err) + } + move(t, pool, tenant, session, first.TurnID, sessions.TurnInProgress, sessions.TurnCancelled) + next := submitMessage(t, service, tenant, session, "next") + for key, original := range map[string]sessions.InputReceipt{"cancel": cancel, "idle-cancel": idle} { + retry, err := submit(ctx, service, tenant, session, key, cancelInput) + if err != nil || !retry.Replayed || retry.TurnID != original.TurnID || retry.Sequence != original.Sequence { + t.Fatalf("cancellation retargeted: %+v, %v", retry, err) + } + } + queued, err := store.GetTurn(ctx, text(tenant), text(session), next.TurnID) + if err != nil || queued.Status != sessions.TurnQueued || !queued.CancelRequestedAt.IsZero() { + t.Fatalf("old cancellation affected later Turn: %+v, %v", queued, err) + } + if _, err := submit(ctx, service, tenant, session, "cancel-queued", cancelInput); err != nil { + t.Fatal(err) + } + stopped, err := store.GetTurn(ctx, text(tenant), text(session), next.TurnID) + if err != nil || stopped.Status != sessions.TurnCancelled || stopped.CompletedAt.IsZero() || !stopped.StartedAt.IsZero() { + t.Fatalf("queued work did not stop: %+v, %v", stopped, err) + } + if _, err := transition(ctx, pool, tenant, session, next.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatalf("cancelled queued work was started: %v", err) + } +} + +// A waiting Turn takes steering input and keeps its start time when it +// resumes; a cancel request keeps it from resuming. +func TestWaitingTurnRetainsInputsAndStartTime(t *testing.T) { + pool := pgtest.Open(t) + _, service := stagingService(t, pool) + tenant, session := newSession(t, pool) + first := submitMessage(t, service, tenant, session, "first") + started := move(t, pool, tenant, session, first.TurnID, sessions.TurnQueued, sessions.TurnInProgress) + move(t, pool, tenant, session, first.TurnID, sessions.TurnInProgress, sessions.TurnWaiting) + steer := submitMessage(t, service, tenant, session, "steer") + if steer.TurnID != first.TurnID { + t.Fatal("waiting input started another Turn") + } + resumed := move(t, pool, tenant, session, first.TurnID, sessions.TurnWaiting, sessions.TurnInProgress) + if !started.StartedAt.Equal(resumed.StartedAt) { + t.Fatal("resume reset the start time") + } + move(t, pool, tenant, session, first.TurnID, sessions.TurnInProgress, sessions.TurnWaiting) + if _, err := submit(t.Context(), service, tenant, session, "cancel", cancelInput); err != nil { + t.Fatal(err) + } + if _, err := transition(t.Context(), pool, tenant, session, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnWaiting, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatalf("cancelling Turn resumed: %v", err) + } + move(t, pool, tenant, session, first.TurnID, sessions.TurnWaiting, sessions.TurnCancelled) +} + +// Rejected input and rejected transitions persist nothing. +func TestTurnInputValidationHasNoSideEffects(t *testing.T) { + pool := pgtest.Open(t) + store, service := stagingService(t, pool) + tenant, session := newSession(t, pool) + ctx := t.Context() + for _, raw := range []json.RawMessage{nil, json.RawMessage(`[]`), json.RawMessage(`null`), json.RawMessage(`{} {}`), json.RawMessage(`{"text":"` + string(make([]byte, 512*1024)) + `"}`)} { + if _, err := submit(ctx, service, tenant, session, "first", sessions.Input{Kind: "message", Payload: raw}); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatalf("invalid input accepted: %v", err) + } + } + cancelled, stop := context.WithCancel(ctx) + stop() + if _, err := submit(cancelled, service, tenant, session, "first", sessions.Input{Kind: "message", Payload: messagePayload}); !errors.Is(err, context.Canceled) { + t.Fatalf("cancelled context: %v", err) + } + first := submitMessage(t, service, tenant, session, "first") + if first.Replayed { + t.Fatal("failed submission persisted a receipt") + } + for _, input := range []sessions.TurnTransition{ + {ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnCompleted}, + {ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnQueued}, + {ExpectedStatus: sessions.TurnCompleted, Status: sessions.TurnInProgress}, + {ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress, Outcome: json.RawMessage(`{"premature":true}`)}, + } { + if _, err := transition(ctx, pool, tenant, session, first.TurnID, input); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatalf("invalid transition accepted: %v", err) + } + } + got, err := store.GetTurn(ctx, text(tenant), text(session), first.TurnID) + if err != nil || got.Status != sessions.TurnQueued || !got.StartedAt.IsZero() { + t.Fatalf("invalid transition changed Turn: %+v, %v", got, err) + } +} diff --git a/services/core/internal/persistence/postgres/sessionpg/turn_reads.go b/services/core/internal/persistence/postgres/sessionpg/turn_reads.go index 9848a68d9..28f96b606 100644 --- a/services/core/internal/persistence/postgres/sessionpg/turn_reads.go +++ b/services/core/internal/persistence/postgres/sessionpg/turn_reads.go @@ -72,7 +72,7 @@ func (s *Store) ListTurns(ctx context.Context, tenant, session, cursor string, l } func (s *Store) ListExecutionWork(ctx context.Context, after string, statuses, connectedDevices []string) ([]sessions.ExecutionWork, error) { - id, devices, err := ExecutionWorkCursor(after, connectedDevices) + id, devices, err := executionWorkCursor(after, connectedDevices) if err != nil { return nil, err } @@ -87,10 +87,10 @@ func (s *Store) ListExecutionWork(ctx context.Context, after string, statuses, c return work, nil } -// ExecutionWorkCursor parses the cursor and connected devices of an execution +// executionWorkCursor parses the cursor and connected devices of an execution // work scan; an empty cursor starts at the first ID. A malformed ID is // sessions.ErrInvalidInput. -func ExecutionWorkCursor(after string, connectedDevices []string) (pgtype.UUID, []pgtype.UUID, error) { +func executionWorkCursor(after string, connectedDevices []string) (pgtype.UUID, []pgtype.UUID, error) { id := pgtype.UUID{Valid: true} var err error if after != "" { diff --git a/services/core/internal/sandbox/providers/configuration_flow_test.go b/services/core/internal/sandbox/providers/configuration_flow_test.go index f400a64af..f043ab26d 100644 --- a/services/core/internal/sandbox/providers/configuration_flow_test.go +++ b/services/core/internal/sandbox/providers/configuration_flow_test.go @@ -149,7 +149,7 @@ func TestAdditionalConfigurationProviderUsesCommonAPIAndStore(t *testing.T) { Execution: &api.Execution{ ExecutorURL: "wss://core.example/api/v1/agent-daemon/ws", SessionAdmission: s, - InputAdmission: s, + InputAdmission: struct{ api.InputAdmission }{}, SessionArchive: struct{ api.SessionArchive }{}, Workspaces: struct{ api.EnvironmentWorkspaces }{}, }, diff --git a/services/core/internal/sessions/artifacts_test.go b/services/core/internal/sessions/artifacts_test.go index a31e5b326..b7a45c855 100644 --- a/services/core/internal/sessions/artifacts_test.go +++ b/services/core/internal/sessions/artifacts_test.go @@ -192,6 +192,11 @@ func (s *fakeArtifactStorage) DeleteSessionArtifact(context.Context, string, str return nil } +func (s *fakeArtifactStorage) WithInputs(context.Context, string, string, func(context.Context, InputTx) error) error { + s.t.Fatal("unexpected call to WithInputs") + return nil +} + func (s *fakeArtifactStorage) CreateDevice(context.Context, string, DeviceRegistration) (ExecutionDevice, error) { s.t.Fatal("unexpected call to CreateDevice") return ExecutionDevice{}, nil diff --git a/services/core/internal/sessions/devices_test.go b/services/core/internal/sessions/devices_test.go index b8297e341..1976c4769 100644 --- a/services/core/internal/sessions/devices_test.go +++ b/services/core/internal/sessions/devices_test.go @@ -45,6 +45,8 @@ type fakeStorage struct { deletion *fakeTx updateMetadata func(encoded string) (Session, error) auditOperation func() error + // inputTx is the transaction WithInputs applies in. + inputTx *fakeInputTx } func (s *fakeStorage) record(name string, set bool, detail ...string) { diff --git a/services/core/internal/sessions/doc.go b/services/core/internal/sessions/doc.go index e86c58aab..069ef6765 100644 --- a/services/core/internal/sessions/doc.go +++ b/services/core/internal/sessions/doc.go @@ -7,12 +7,13 @@ // The Session writes that several operations share are procedures here, over // the transaction interfaces declared beside them: cancelling work, failing // and terminating an Environment, tracking input activity, the admission -// gates, moving a Turn's status, admitting a function result, reading a Turn's -// required actions, appending to a Turn's execution journal and projecting its -// observations and admitted inputs into Items, Turn usage and Subagents. Service runs the pooled -// Session use cases, such as staging a Turn's Artifacts, over Storage, and -// Reader declares the Session reads. ExecutionOperations runs the Session -// writes only the execution owner makes, such as moving and completing Turns, -// recording function calls and their application receipts and journaling a -// Turn's observations, over the lease-bound ExecutionStorage. +// gates, admitting an input batch, moving a Turn's status, admitting a function +// result, reading a Turn's required actions, appending to a Turn's execution +// journal and projecting its observations and admitted inputs into Items, Turn +// usage and Subagents. Service runs the pooled Session use cases, such as +// staging a Turn's Artifacts, over Storage, and Reader declares the Session +// reads. ExecutionOperations runs the Session writes only the execution owner +// makes, such as moving and completing Turns, recording function calls and +// their application receipts and journaling a Turn's observations, over the +// lease-bound ExecutionStorage. package sessions diff --git a/services/core/internal/sessions/environment.go b/services/core/internal/sessions/environment.go index 745ecc791..107aa894c 100644 --- a/services/core/internal/sessions/environment.go +++ b/services/core/internal/sessions/environment.go @@ -98,6 +98,8 @@ const ( type EnvironmentInputReservation struct { ID string SessionID string + // Key is the idempotency key the reservation's batch is admitted under. + Key string State string IsInitial bool Inputs []Input diff --git a/services/core/internal/sessions/environment_inputs.go b/services/core/internal/sessions/environment_inputs.go new file mode 100644 index 000000000..c0b02bee8 --- /dev/null +++ b/services/core/internal/sessions/environment_inputs.go @@ -0,0 +1,159 @@ +package sessions + +import ( + "context" + "errors" +) + +// checkEnvironmentAcceptsInput admits new input only while the Session's +// Environment is live: a failed hosted Environment is +// ErrHostedEnvironmentFailed, and any other failed or expired Environment +// ErrEnvironmentUnavailable. +func checkEnvironmentAcceptsInput(environment Environment) error { + if environment.Status == "failed" { + if kind, err := EnvironmentType(environment.Configuration); err == nil && kind == "openai_hosted" { + return ErrHostedEnvironmentFailed + } + } + if environment.Status == "failed" || environment.Status == "expired" { + return ErrEnvironmentUnavailable + } + return nil +} + +// joinsActiveTurn reports whether reserved Environment input joins the +// Session's active Turn at once instead of waiting for its own Turn: it does +// while the Turn has not started capturing its Artifacts, whose sealed native +// input takes nothing more. +func joinsActiveTurn(turn Turn, active bool) bool { + return active && !turn.ArtifactCaptureStarted +} + +// ReserveEnvironmentInput admits a message batch for a Session whose +// Environment runs it. Under the Session lock it joins the active Turn at once, +// unless that Turn captures its Artifacts, or otherwise reserves the batch +// until a deadline, for the execution owner to promote into a new Turn. A +// retry returns the reservation under its key, expired once its deadline has +// passed, or the receipts the batch was admitted with. The checks of +// SubmitInputs apply, and a failed or expired Environment rejects new input. +func (s *Service) ReserveEnvironmentInput(ctx context.Context, tenant, session, key string, inputs []Input) (EnvironmentInputReservation, error) { + if err := ValidateInputKey(key); err != nil { + return EnvironmentInputReservation{}, err + } + batch, encoded, err := ValidateMessageInputs(inputs) + if err != nil { + return EnvironmentInputReservation{}, err + } + var result EnvironmentInputReservation + err = s.storage.WithInputs(ctx, tenant, session, func(ctx context.Context, tx InputTx) error { + return TrackInputActivity(ctx, tx, func(ctx context.Context) error { + previous, used, err := tx.FindInputReservation(ctx, key, encoded) + if err != nil { + return err + } + replay, err := decideReplay(used, previous != nil) + if err != nil { + return err + } + if replay { + if result, err = expireReservation(ctx, tx, *previous); err != nil { + return err + } + if result.State == EnvironmentInputPending || result.State == EnvironmentInputAdmitted { + return tx.RecordInputAudit(ctx) + } + return nil + } + receipts, matches, err := tx.LoadInputBatch(ctx, key, encoded) + if err != nil { + return err + } + if replay, err = decideReplay(len(receipts) > 0, matches); err != nil { + return err + } + if replay { + // Earlier direct admission has receipts, but never had a reservation or deadline. + result = EnvironmentInputReservation{SessionID: session, State: EnvironmentInputAdmitted, Receipts: receipts} + return tx.RecordInputAudit(ctx) + } + if err := CheckFileWriteGate(ctx, tx); err != nil { + return err + } + environment, err := tx.LoadEnvironment(ctx) + if errors.Is(err, ErrNotFound) { + return ErrInvalidInput + } + if err != nil { + return err + } + if err := checkEnvironmentAcceptsInput(environment); err != nil { + return err + } + gate, err := tx.LoadInputGate(ctx, key, encoded) + if err != nil { + return err + } + if err := checkInputGate(gate); err != nil { + return err + } + turn, active, err := tx.LoadActiveTurn(ctx) + if err != nil { + return err + } + if joinsActiveTurn(turn, active) { + receipts, err := AdmitInputs(ctx, tx, key, batch) + if err != nil { + return err + } + result = EnvironmentInputReservation{SessionID: session, State: EnvironmentInputAdmitted, Receipts: receipts} + return tx.RecordInputAudit(ctx) + } + if result, err = tx.CreateInputReservation(ctx, key, encoded); err != nil { + return err + } + return tx.RecordInputAudit(ctx) + }) + }) + if err != nil { + return EnvironmentInputReservation{}, err + } + return result, nil +} + +// ExpireEnvironmentInput returns the Session's Environment input reservation, +// expired first when it is pending and its deadline has passed by the +// database clock. A malformed reservation ID is ErrInvalidInput, and a missing +// reservation ErrNotFound. +func (s *Service) ExpireEnvironmentInput(ctx context.Context, tenant, session, reservation string) (EnvironmentInputReservation, error) { + if !validID(reservation) { + return EnvironmentInputReservation{}, ErrInvalidInput + } + var result EnvironmentInputReservation + err := s.storage.WithInputs(ctx, tenant, session, func(ctx context.Context, tx InputTx) error { + return TrackInputActivity(ctx, tx, func(ctx context.Context) error { + current, err := tx.LoadInputReservation(ctx, reservation) + if err != nil { + return err + } + result, err = expireReservation(ctx, tx, current) + return err + }) + }) + if err != nil { + return EnvironmentInputReservation{}, err + } + return result, nil +} + +// expireReservation expires a pending reservation whose deadline has passed +// and returns it as the transaction has left it. A settled one is returned as +// it is: its outcome is a successful result, so settling never rolls back. +func expireReservation(ctx context.Context, tx InputTx, reservation EnvironmentInputReservation) (EnvironmentInputReservation, error) { + if reservation.State != EnvironmentInputPending { + return reservation, nil + } + if err := tx.ExpireInputReservation(ctx, reservation.ID); err != nil { + return EnvironmentInputReservation{}, err + } + return tx.LoadInputReservation(ctx, reservation.ID) +} diff --git a/services/core/internal/sessions/environment_inputs_test.go b/services/core/internal/sessions/environment_inputs_test.go new file mode 100644 index 000000000..f49fd2aaf --- /dev/null +++ b/services/core/internal/sessions/environment_inputs_test.go @@ -0,0 +1,165 @@ +package sessions + +import ( + "encoding/json" + "errors" + "reflect" + "testing" +) + +const testReservation = "6a1c9e2b-3d4f-4a5b-8c6d-7e8f9a0b1c2d" + +// reservations is a fake LoadInputReservation that reads states one call at +// a time, then the last again. +func reservations(states ...EnvironmentInputReservation) func() (EnvironmentInputReservation, error) { + return func() (EnvironmentInputReservation, error) { + state := states[0] + if len(states) > 1 { + states = states[1:] + } + return state, nil + } +} + +func reservationIn(state string) EnvironmentInputReservation { + return EnvironmentInputReservation{ID: testReservation, SessionID: testSession, Key: inputKey, State: state, Inputs: []Input{messageInput("hi")}} +} + +func TestCheckEnvironmentAcceptsInput(t *testing.T) { + hosted := json.RawMessage(`{"type":"openai_hosted"}`) + for _, test := range []struct { + environment Environment + want error + }{ + {selfHosted, nil}, + {Environment{Status: "failed", Configuration: hosted}, ErrHostedEnvironmentFailed}, + {Environment{Status: "failed", Configuration: selfHosted.Configuration}, ErrEnvironmentUnavailable}, + {Environment{Status: "expired", Configuration: hosted}, ErrEnvironmentUnavailable}, + } { + if err := checkEnvironmentAcceptsInput(test.environment); err != test.want { + t.Errorf("%s %s: %v", test.environment.Status, test.environment.Configuration, err) + } + } +} + +func TestJoinsActiveTurn(t *testing.T) { + capturing := turnWith(TurnInProgress) + capturing.ArtifactCaptureStarted = true + if joinsActiveTurn(Turn{}, false) || !joinsActiveTurn(turnWith(TurnInProgress), true) || joinsActiveTurn(capturing, true) { + t.Fatal("joins") + } +} + +func TestReserveEnvironmentInput(t *testing.T) { + const batch = `[{"kind":"message","payload":{"text":"hi"}}]` + find := "FindInputReservation request " + batch + reserve := func(t *testing.T, tx *fakeInputTx) (EnvironmentInputReservation, error) { + tx.loadEnvironmentInput = inputs() + return deviceService(t, &fakeStorage{t: t, inputTx: tx}).ReserveEnvironmentInput(t.Context(), testTenant, testSession, inputKey, []Input{messageInput("hi")}) + } + // fresh is a transaction where the key holds no batch yet. + fresh := func(t *testing.T) *fakeInputTx { + tx := newInputTx(t) + tx.findInputReservation = func() (*EnvironmentInputReservation, bool, error) { return nil, false, nil } + tx.loadInputBatch = func() ([]InputReceipt, bool, error) { return nil, true, nil } + tx.loadPendingFileWrite, tx.loadEnvironment = returns(false), returns(selfHosted) + return tx + } + checked := []string{"LoadEnvironmentInput", find, "LoadInputBatch request " + batch, "LoadPendingFileWrite", "LoadEnvironment"} + t.Run("a retry returns its reservation", func(t *testing.T) { + for _, test := range []struct { + name string + current EnvironmentInputReservation + audit []string + }{ + {"pending", reservationIn(EnvironmentInputPending), []string{"RecordInputAudit"}}, + {"expired at its deadline", reservationIn(EnvironmentInputExpired), nil}, + } { + t.Run(test.name, func(t *testing.T) { + previous := reservationIn(EnvironmentInputPending) + tx := newInputTx(t) + tx.findInputReservation = func() (*EnvironmentInputReservation, bool, error) { return &previous, true, nil } + tx.expireInputReservation, tx.loadInputReservation, tx.recordInputAudit = done, returns(test.current), done + result, err := reserve(t, tx) + if err != nil || !reflect.DeepEqual(result, test.current) { + t.Fatalf("reservation %+v, %v", result, err) + } + calls := append([]string{"LoadEnvironmentInput", find, "ExpireInputReservation " + testReservation, "LoadInputReservation " + testReservation}, test.audit...) + assertCalls(t, tx.fakeTx, append(calls, "LoadEnvironmentInput")...) + }) + } + }) + t.Run("a retry of a direct admission returns its receipts", func(t *testing.T) { + previous := []InputReceipt{{Sequence: 3, TurnID: testTurn, Replayed: true}} + tx := newInputTx(t) + tx.findInputReservation = func() (*EnvironmentInputReservation, bool, error) { return nil, false, nil } + tx.loadInputBatch = func() ([]InputReceipt, bool, error) { return previous, true, nil } + tx.recordInputAudit = done + result, err := reserve(t, tx) + if err != nil || !reflect.DeepEqual(result, EnvironmentInputReservation{SessionID: testSession, State: EnvironmentInputAdmitted, Receipts: previous}) { + t.Fatalf("reservation %+v, %v", result, err) + } + assertCalls(t, tx.fakeTx, "LoadEnvironmentInput", find, "LoadInputBatch request "+batch, "RecordInputAudit", "LoadEnvironmentInput") + }) + t.Run("an Environment that no longer runs rejects input", func(t *testing.T) { + for _, test := range []struct { + name string + load func() (Environment, error) + want error + }{ + {"failed hosted", returns(Environment{Status: "failed", Configuration: json.RawMessage(`{"type":"openai_hosted"}`)}), ErrHostedEnvironmentFailed}, + {"expired", returns(Environment{Status: "expired", Configuration: selfHosted.Configuration}), ErrEnvironmentUnavailable}, + {"missing", func() (Environment, error) { return Environment{}, ErrNotFound }, ErrInvalidInput}, + } { + t.Run(test.name, func(t *testing.T) { + tx := fresh(t) + tx.loadEnvironment = test.load + if _, err := reserve(t, tx); err != test.want { + t.Fatal(err) + } + assertCalls(t, tx.fakeTx, checked...) + }) + } + }) + t.Run("input joins the active Turn", func(t *testing.T) { + running := turnWith(TurnInProgress) + tx := fresh(t) + tx.loadInputGate, tx.loadActiveTurn, tx.createTurnInput, tx.loadInputSource, tx.recordInputAudit = returns(InputGate{Matches: true}), activeTurn(&running), sequences(4), unprojected, done + result, err := reserve(t, tx) + if err != nil || !reflect.DeepEqual(result, EnvironmentInputReservation{SessionID: testSession, State: EnvironmentInputAdmitted, Receipts: []InputReceipt{{Sequence: 4, TurnID: testTurn}}}) { + t.Fatalf("reservation %+v, %v", result, err) + } + assertCalls(t, tx.fakeTx, append(checked, "LoadInputGate request "+batch, "LoadActiveTurn", + "LoadActiveTurn", "CreateTurnInput "+testTurn+" request 0 message", "LoadInputSource 4", "RecordInputAudit", "LoadEnvironmentInput")...) + }) + t.Run("input waits in a reservation on an idle Session", func(t *testing.T) { + tx := fresh(t) + tx.loadInputGate, tx.loadActiveTurn = returns(InputGate{Matches: true}), activeTurn(nil) + tx.createInputReservation, tx.recordInputAudit = returns(reservationIn(EnvironmentInputPending)), done + result, err := reserve(t, tx) + if err != nil || !reflect.DeepEqual(result, reservationIn(EnvironmentInputPending)) { + t.Fatalf("reservation %+v, %v", result, err) + } + assertCalls(t, tx.fakeTx, append(checked, "LoadInputGate request "+batch, "LoadActiveTurn", "CreateInputReservation request "+batch, "RecordInputAudit", "LoadEnvironmentInput")...) + }) + t.Run("a cancellation touches no storage", func(t *testing.T) { + if _, err := deviceService(t, &fakeStorage{t: t}).ReserveEnvironmentInput(t.Context(), testTenant, testSession, inputKey, []Input{cancelInput}); !errors.Is(err, ErrInvalidInput) { + t.Fatal(err) + } + }) +} + +func TestExpireEnvironmentInput(t *testing.T) { + tx := newInputTx(t) + tx.loadEnvironmentInput, tx.expireInputReservation = inputs(), done + tx.loadInputReservation = reservations(reservationIn(EnvironmentInputPending), reservationIn(EnvironmentInputExpired)) + service := deviceService(t, &fakeStorage{t: t, inputTx: tx}) + result, err := service.ExpireEnvironmentInput(t.Context(), testTenant, testSession, testReservation) + if err != nil || result.State != EnvironmentInputExpired { + t.Fatalf("reservation %+v, %v", result, err) + } + assertCalls(t, tx.fakeTx, "LoadEnvironmentInput", "LoadInputReservation "+testReservation, "ExpireInputReservation "+testReservation, "LoadInputReservation "+testReservation, "LoadEnvironmentInput") + if _, err := service.ExpireEnvironmentInput(t.Context(), testTenant, testSession, "malformed"); !errors.Is(err, ErrInvalidInput) { + t.Fatal(err) + } +} diff --git a/services/core/internal/sessions/execution.go b/services/core/internal/sessions/execution.go index db5eaaf13..c07c25573 100644 --- a/services/core/internal/sessions/execution.go +++ b/services/core/internal/sessions/execution.go @@ -20,5 +20,6 @@ func NewExecutionOperations(storage ExecutionStorage) (*ExecutionOperations, err type ExecutionStorage interface { EnvironmentExecution FunctionExecution + InputExecution TurnExecution } diff --git a/services/core/internal/sessions/execution_functions_test.go b/services/core/internal/sessions/execution_functions_test.go index cc8ca606a..03cdbb09e 100644 --- a/services/core/internal/sessions/execution_functions_test.go +++ b/services/core/internal/sessions/execution_functions_test.go @@ -15,9 +15,11 @@ import ( // their call on tx, then run apply with tx and locked; without tx they fail // the test. type fakeExecutionStorage struct { - t *testing.T - withFunctionTurn func(ctx context.Context, tenant, session, turn string, apply func(context.Context, FunctionTx, Turn) error) error - withTurns func(ctx context.Context, tenant, session string, apply func(context.Context, TurnTx) error) error + t *testing.T + withFunctionTurn func(ctx context.Context, tenant, session, turn string, apply func(context.Context, FunctionTx, Turn) error) error + withTurns func(ctx context.Context, tenant, session string, apply func(context.Context, TurnTx) error) error + withInputs func(ctx context.Context, tenant, session string, apply func(context.Context, InputTx) error) error + withDueInputReservations func(ctx context.Context, apply func(context.Context, InputTx, string) error) error tx *fakeTx locked LockedSession diff --git a/services/core/internal/sessions/execution_inputs.go b/services/core/internal/sessions/execution_inputs.go new file mode 100644 index 000000000..111195033 --- /dev/null +++ b/services/core/internal/sessions/execution_inputs.go @@ -0,0 +1,103 @@ +package sessions + +import "context" + +// InputExecution is the lease-bound storage of Environment input settlement. +type InputExecution interface { + // WithInputs runs apply in one transaction on the execution lease, under + // the lock of the tenant's Session, and commits only when apply succeeds. + // A malformed tenant is ErrInvalidInput; a malformed, missing or deleted + // Session is ErrNotFound. + WithInputs(ctx context.Context, tenant, session string, apply func(context.Context, InputTx) error) error + // WithDueInputReservations runs apply in one transaction on the execution + // lease for each of at most 32 pending Environment input reservations + // whose deadline has passed, earliest deadline first, of live Sessions + // whose lock no other transaction holds, with that Session's transaction + // and the reservation's ID. It commits only when every apply succeeds. + WithDueInputReservations(ctx context.Context, apply func(context.Context, InputTx, string) error) error +} + +// PromoteEnvironmentInput admits a pending reservation's batch into a new Turn +// and claims that Turn as in progress for the execution owner's retained +// native preparation. The reservation settles as admitted with the batch's +// receipts. A reservation that already settled, or expires now because its +// deadline has passed, is returned as it is. The Session must be idle with no +// pending file write; otherwise it is ErrTurnConflict and the reservation +// stays pending. A malformed reservation ID is ErrInvalidInput, and a missing +// reservation ErrNotFound. +func (o *ExecutionOperations) PromoteEnvironmentInput(ctx context.Context, tenant, session, reservation string) (EnvironmentInputReservation, error) { + if !validID(reservation) { + return EnvironmentInputReservation{}, ErrInvalidInput + } + var result EnvironmentInputReservation + err := o.storage.WithInputs(ctx, tenant, session, func(ctx context.Context, tx InputTx) error { + return TrackInputActivity(ctx, tx, func(ctx context.Context) error { + current, err := tx.LoadInputReservation(ctx, reservation) + if err != nil { + return err + } + if result, err = expireReservation(ctx, tx, current); err != nil || result.State != EnvironmentInputPending { + return err + } + if err := CheckInputStart(ctx, tx); err != nil { + return err + } + if result.Receipts, err = AdmitInputs(ctx, tx, result.Key, result.Inputs); err != nil { + return err + } + settled, err := tx.AdmitInputReservation(ctx, reservation) + if err != nil { + return err + } + start := TurnTransition{ExpectedStatus: TurnQueued, Status: TurnInProgress} + if _, err := TransitionTurn(ctx, tx, result.Receipts[0].TurnID, start); err != nil { + return err + } + result.State, result.SettledAt = EnvironmentInputAdmitted, &settled + return nil + }) + }) + if err != nil { + return EnvironmentInputReservation{}, err + } + return result, nil +} + +// FailEnvironmentInput settles a pending reservation as failed before +// admission with code, model_provider_required or runtime_preparation_failed, +// once execution confirmed the failure. A reservation that already settled, +// including one cancelled meanwhile, keeps its outcome, and one whose deadline +// has passed expires instead. Any other code and a malformed reservation ID +// are ErrInvalidInput. +func (o *ExecutionOperations) FailEnvironmentInput(ctx context.Context, tenant, session, reservation, code string) error { + if (code != "model_provider_required" && code != "runtime_preparation_failed") || !validID(reservation) { + return ErrInvalidInput + } + return o.storage.WithInputs(ctx, tenant, session, func(ctx context.Context, tx InputTx) error { + return TrackInputActivity(ctx, tx, func(ctx context.Context) error { + if err := tx.ExpireInputReservation(ctx, reservation); err != nil { + return err + } + return tx.FailInputReservation(ctx, reservation, code) + }) + }) +} + +// ExpireEnvironmentInputs expires one bounded batch of pending reservations +// whose deadline has passed, reports the input activity each changes, creates +// no Turn history and returns how many it expired. It is cross-Session +// maintenance only the execution owner runs. +func (o *ExecutionOperations) ExpireEnvironmentInputs(ctx context.Context) (int64, error) { + var expired int64 + err := o.storage.WithDueInputReservations(ctx, func(ctx context.Context, tx InputTx, reservation string) error { + if err := TrackInputActivity(ctx, tx, func(ctx context.Context) error { return tx.ExpireInputReservation(ctx, reservation) }); err != nil { + return err + } + expired++ + return nil + }) + if err != nil { + return 0, err + } + return expired, nil +} diff --git a/services/core/internal/sessions/execution_inputs_test.go b/services/core/internal/sessions/execution_inputs_test.go new file mode 100644 index 000000000..64a10a1bd --- /dev/null +++ b/services/core/internal/sessions/execution_inputs_test.go @@ -0,0 +1,133 @@ +package sessions + +import ( + "context" + "encoding/json" + "errors" + "reflect" + "testing" + "time" +) + +func (f *fakeExecutionStorage) WithInputs(ctx context.Context, tenant, session string, apply func(context.Context, InputTx) error) error { + f.t.Helper() + if f.withInputs == nil { + f.t.Fatal("unexpected call to WithInputs") + } + return f.withInputs(ctx, tenant, session, apply) +} + +func (f *fakeExecutionStorage) WithDueInputReservations(ctx context.Context, apply func(context.Context, InputTx, string) error) error { + f.t.Helper() + if f.withDueInputReservations == nil { + f.t.Fatal("unexpected call to WithDueInputReservations") + } + return f.withDueInputReservations(ctx, apply) +} + +// inputOperations serves the tenant's Session through tx. +func inputOperations(t *testing.T, tx *fakeInputTx) *ExecutionOperations { + t.Helper() + tx.loadEnvironmentInput = inputs() + operations, err := NewExecutionOperations(&fakeExecutionStorage{t: t, withInputs: func(ctx context.Context, tenant, session string, apply func(context.Context, InputTx) error) error { + if tenant != testTenant || session != testSession { + t.Fatalf("Session %s %s", tenant, session) + } + return apply(ctx, tx) + }}) + if err != nil { + t.Fatal(err) + } + return operations +} + +func TestPromoteEnvironmentInput(t *testing.T) { + expire := []string{"LoadEnvironmentInput", "LoadInputReservation " + testReservation, "ExpireInputReservation " + testReservation, "LoadInputReservation " + testReservation} + promote := func(t *testing.T, tx *fakeInputTx) (EnvironmentInputReservation, error) { + return inputOperations(t, tx).PromoteEnvironmentInput(t.Context(), testTenant, testSession, testReservation) + } + t.Run("a reservation past its deadline expires", func(t *testing.T) { + tx := newInputTx(t) + tx.loadInputReservation, tx.expireInputReservation = reservations(reservationIn(EnvironmentInputPending), reservationIn(EnvironmentInputExpired)), done + result, err := promote(t, tx) + if err != nil || !reflect.DeepEqual(result, reservationIn(EnvironmentInputExpired)) { + t.Fatalf("reservation %+v, %v", result, err) + } + assertCalls(t, tx.fakeTx, append(expire, "LoadEnvironmentInput")...) + }) + t.Run("a pending reservation starts its Turn", func(t *testing.T) { + settled := time.Unix(1700000002, 0) + tx := newInputTx(t) + tx.loadInputReservation, tx.expireInputReservation = returns(reservationIn(EnvironmentInputPending)), done + tx.loadPendingFileWrite, tx.loadActiveTurn = returns(false), activeTurn(nil) + tx.createTurn, tx.createTurnInput, tx.loadInputSource = returns(turnWith(TurnQueued)), sequences(5), unprojected + tx.loadUsage, tx.appendChanges = returns(json.RawMessage(`{}`)), collect(new([]SessionChange)) + tx.admitInputReservation, tx.loadTurn, tx.loadComputeSuspension = returns(settled), returns(turnWith(TurnQueued)), returns(false) + tx.applyTurnStatus = func(TurnStatusChange) (Turn, error) { return turnWith(TurnInProgress), nil } + result, err := promote(t, tx) + want := reservationIn(EnvironmentInputAdmitted) + want.SettledAt, want.Receipts = &settled, []InputReceipt{{Sequence: 5, TurnID: testTurn}} + if err != nil || !reflect.DeepEqual(result, want) { + t.Fatalf("reservation %+v, %v", result, err) + } + assertCalls(t, tx.fakeTx, append(expire, "LoadPendingFileWrite", "LoadActiveTurn", + "LoadActiveTurn", "CreateTurn", "AppendChanges agent.session.turn.created", "CreateTurnInput "+testTurn+" request 0 message", "LoadInputSource 5", "LoadUsage", "AppendChanges agent.session.in_progress", + "AdmitInputReservation "+testReservation, + "LoadTurn "+testTurn, "LoadComputeSuspension", "ApplyTurnStatus "+testTurn+" queued in_progress {}", "AppendChanges agent.session.turn.in_progress", + "LoadEnvironmentInput")...) + }) + t.Run("a busy Session keeps the reservation pending", func(t *testing.T) { + running := turnWith(TurnInProgress) + tx := newInputTx(t) + tx.loadInputReservation, tx.expireInputReservation = returns(reservationIn(EnvironmentInputPending)), done + tx.loadPendingFileWrite, tx.loadActiveTurn = returns(false), activeTurn(&running) + if _, err := promote(t, tx); !errors.Is(err, ErrTurnConflict) { + t.Fatal(err) + } + assertCalls(t, tx.fakeTx, append(expire, "LoadPendingFileWrite", "LoadActiveTurn")...) + }) + if _, err := unusedStorage(t).PromoteEnvironmentInput(t.Context(), testTenant, testSession, "malformed"); !errors.Is(err, ErrInvalidInput) { + t.Fatal(err) + } +} + +func TestFailEnvironmentInput(t *testing.T) { + tx := newInputTx(t) + tx.expireInputReservation, tx.failInputReservation = done, done + if err := inputOperations(t, tx).FailEnvironmentInput(t.Context(), testTenant, testSession, testReservation, "runtime_preparation_failed"); err != nil { + t.Fatal(err) + } + assertCalls(t, tx.fakeTx, "LoadEnvironmentInput", "ExpireInputReservation "+testReservation, "FailInputReservation "+testReservation+" runtime_preparation_failed", "LoadEnvironmentInput") + for reservation, code := range map[string]string{testReservation: "other", "malformed": "model_provider_required"} { + if err := unusedStorage(t).FailEnvironmentInput(t.Context(), testTenant, testSession, reservation, code); !errors.Is(err, ErrInvalidInput) { + t.Fatalf("%s %s: %v", reservation, code, err) + } + } +} + +func TestExpireEnvironmentInputs(t *testing.T) { + due := []string{"6a1c9e2b-0000-4a5b-8c6d-7e8f9a0b1c2d", "6a1c9e2b-1111-4a5b-8c6d-7e8f9a0b1c2d"} + expire := func(t *testing.T, tx *fakeInputTx, failure error) (int64, error) { + tx.loadEnvironmentInput, tx.expireInputReservation = inputs(), done + operations, err := NewExecutionOperations(&fakeExecutionStorage{t: t, withDueInputReservations: func(ctx context.Context, apply func(context.Context, InputTx, string) error) error { + for _, reservation := range due { + if err := apply(ctx, tx, reservation); err != nil { + return err + } + } + return failure + }}) + if err != nil { + t.Fatal(err) + } + return operations.ExpireEnvironmentInputs(t.Context()) + } + tx := newInputTx(t) + if expired, err := expire(t, tx, nil); err != nil || expired != 2 { + t.Fatalf("expired %d, %v", expired, err) + } + assertCalls(t, tx.fakeTx, "LoadEnvironmentInput", "ExpireInputReservation "+due[0], "LoadEnvironmentInput", "LoadEnvironmentInput", "ExpireInputReservation "+due[1], "LoadEnvironmentInput") + if expired, err := expire(t, newInputTx(t), errStorage); !errors.Is(err, errStorage) || expired != 0 { + t.Fatalf("expired %d, %v", expired, err) + } +} diff --git a/services/core/internal/sessions/inputs.go b/services/core/internal/sessions/inputs.go index e041795b5..264cbb106 100644 --- a/services/core/internal/sessions/inputs.go +++ b/services/core/internal/sessions/inputs.go @@ -1,10 +1,14 @@ package sessions import ( + "context" "encoding/json" "fmt" + "slices" "strings" "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" ) // Input is a validated execution command, not an upstream wire type. @@ -35,6 +39,334 @@ func ValidateInputKey(key string) error { return nil } +// ValidateInputs checks and normalizes an input batch: 1..64 messages, +// cancellations and function results whose nonempty JSON object payloads total +// at most 512 KiB, an empty cancellation payload and a well-formed function +// result. It returns the normalized batch and its encoding, the batch's retry +// identity. +func ValidateInputs(inputs []Input) ([]Input, json.RawMessage, error) { + if len(inputs) == 0 || len(inputs) > 64 { + return nil, nil, fmt.Errorf("%w: input batch must contain 1..64 events", ErrInvalidInput) + } + batch := make([]Input, len(inputs)) + size := 0 + for i, input := range inputs { + size += len(input.Payload) + if size > 512*1024 || len(input.Payload) == 0 || (input.Kind != "message" && input.Kind != "cancel" && input.Kind != "tool_result") { + return nil, nil, fmt.Errorf("%w: input payloads must be nonempty and total at most 512 KiB", ErrInvalidInput) + } + payload, err := jsonobject.Normalize(input.Payload) + if err != nil { + return nil, nil, fmt.Errorf("%w: %w", ErrInvalidInput, err) + } + if input.Kind == "cancel" && string(payload) != "{}" { + return nil, nil, fmt.Errorf("%w: cancel payload must be empty", ErrInvalidInput) + } + if input.Kind == "tool_result" { + if _, err := ParseFunctionResultInput(payload); err != nil { + return nil, nil, err + } + } + batch[i] = Input{Kind: input.Kind, Payload: payload} + } + encoded, err := json.Marshal(batch) + return batch, encoded, err +} + +// ValidateMessageInputs validates a batch as ValidateInputs does and also +// requires that it holds only messages, as Environment input reservations and +// Session initial input do. +func ValidateMessageInputs(inputs []Input) ([]Input, json.RawMessage, error) { + for _, input := range inputs { + if input.Kind != "message" { + return nil, nil, ErrInvalidInput + } + } + return ValidateInputs(inputs) +} + +// decideReplay decides a request whose idempotency key may already hold a +// batch: an unused key admits the request, a key holding the same batch +// replays what it admitted, and a key holding a different batch is +// ErrIdempotencyConflict. The whole batch is the retry identity, so a replay +// never re-evaluates its inputs, such as a cancellation's target. +func decideReplay(used, matches bool) (bool, error) { + if !used { + return false, nil + } + if !matches { + return false, ErrIdempotencyConflict + } + return true, nil +} + +// InputGate is what the Session's Environment input reservations show about a +// new input batch. +type InputGate struct { + // Matches reports that the Session has no reservation under the batch's + // idempotency key, or one with the same batch. + Matches bool + // Blocked reports that a reservation is pending or one holds the batch's + // idempotency key. + Blocked bool +} + +// checkInputGate admits a new input batch only while no Environment input +// reservation is pending or holds its key: a reservation holding the key with +// a different batch is ErrIdempotencyConflict, and any other blocking +// reservation ErrInputPending. +func checkInputGate(gate InputGate) error { + if !gate.Matches { + return ErrIdempotencyConflict + } + if gate.Blocked { + return ErrInputPending + } + return nil +} + +// inputPlacement is where admission puts a message or cancellation. +type inputPlacement int + +const ( + // steersTurn adds a message to the active Turn. + steersTurn inputPlacement = iota + // cancelsTurn adds a cancellation to the active Turn and requests it. + cancelsTurn + // startsTurn starts a new Turn with a message on an idle Session. + startsTurn + // keepsIdentity records a cancellation on an idle Session without a Turn, + // which keeps only its retry identity. + keepsIdentity +) + +// placeInput decides where a message or cancellation goes, given whether the +// Session has an active Turn. +func placeInput(kind string, active bool) inputPlacement { + switch { + case active && kind == "cancel": + return cancelsTurn + case active: + return steersTurn + case kind == "message": + return startsTurn + } + return keepsIdentity +} + +// InputAdmissionTx is the Session transaction AdmitInputs runs in. +type InputAdmissionTx interface { + FunctionResultTx + TurnCancellationTx + InputProjectionTx + ActiveTurnTx + // CreateTurn creates a queued Turn in the Session and returns it as + // stored. + CreateTurn(ctx context.Context) (Turn, error) + // CreateTurnInput records input at position of the batch admitted under + // key, joined to turn, or to no Turn when turn is empty, and returns the + // Session sequence it allocates. + CreateTurnInput(ctx context.Context, turn, key string, position int32, input Input) (int64, error) +} + +// AdmitInputs admits the validated batch under key, input by input in batch +// order, and returns their receipts. A function result joins the Turn whose +// call it answers. A message steers the Session's active Turn or, on an idle +// Session, starts a new Turn, which publishes turn.created, then the message's +// Items, then the Session activity. A cancellation requests the cancellation +// of the active Turn; on an idle Session it joins no Turn and keeps only its +// retry identity. +func AdmitInputs(ctx context.Context, tx InputAdmissionTx, key string, batch []Input) ([]InputReceipt, error) { + receipts := make([]InputReceipt, 0, len(batch)) + for position, input := range batch { + receipt, err := admitInput(ctx, tx, key, int32(position), input) + if err != nil { + return nil, err + } + receipts = append(receipts, receipt) + } + return receipts, nil +} + +// admitInput admits the input at position of the batch under key. +func admitInput(ctx context.Context, tx InputAdmissionTx, key string, position int32, input Input) (InputReceipt, error) { + if input.Kind == "tool_result" { + result, err := ParseFunctionResultInput(input.Payload) + if err != nil { + return InputReceipt{}, err + } + turn, err := AdmitFunctionResult(ctx, tx, result) + if err != nil { + return InputReceipt{}, err + } + sequence, err := tx.CreateTurnInput(ctx, turn.ID, key, position, input) + if err != nil { + return InputReceipt{}, err + } + return InputReceipt{Sequence: sequence, TurnID: turn.ID}, nil + } + turn, active, err := tx.LoadActiveTurn(ctx) + if err != nil { + return InputReceipt{}, err + } + placement := placeInput(input.Kind, active) + if placement == startsTurn { + if turn, err = tx.CreateTurn(ctx); err != nil { + return InputReceipt{}, err + } + if err := tx.AppendChanges(ctx, TurnChanges(turn, true)...); err != nil { + return InputReceipt{}, err + } + } + sequence, err := tx.CreateTurnInput(ctx, turn.ID, key, position, input) + if err != nil { + return InputReceipt{}, err + } + if placement == cancelsTurn { + if err := CancelTurn(ctx, tx, turn); err != nil { + return InputReceipt{}, err + } + } + if err := ProjectInput(ctx, tx, sequence); err != nil { + return InputReceipt{}, err + } + if placement == startsTurn { + usage, err := tx.LoadUsage(ctx) + if err != nil { + return InputReceipt{}, err + } + if err := tx.AppendChanges(ctx, ActivityChange(turn, usage, nil)); err != nil { + return InputReceipt{}, err + } + } + return InputReceipt{Sequence: sequence, TurnID: turn.ID}, nil +} + +// InputTx is the Session transaction of the input use cases. +type InputTx interface { + InputAdmissionTx + InputActivityTx + InputStartTx + TurnTransitionTx + // LoadInputBatch reads the receipts of the inputs the Session admitted + // under key, in batch order and marked replayed, none when it admitted + // none, and reports whether they are batch. + LoadInputBatch(ctx context.Context, key string, batch json.RawMessage) ([]InputReceipt, bool, error) + // LoadInputGate reads what the Session's Environment input reservations + // show about batch under key. + LoadInputGate(ctx context.Context, key string, batch json.RawMessage) (InputGate, error) + // LoadEnvironment reads the Session's Environment; a Session without one + // is ErrNotFound. + LoadEnvironment(ctx context.Context) (Environment, error) + // FindInputReservation reads the Session's Environment input reservation + // under key when its batch is batch, nil otherwise, and reports whether + // the Session reserved key at all. + FindInputReservation(ctx context.Context, key string, batch json.RawMessage) (*EnvironmentInputReservation, bool, error) + // LoadInputReservation reads one of the Session's Environment input + // reservations; a missing one is ErrNotFound. An admitted reservation + // carries the receipts of its batch. + LoadInputReservation(ctx context.Context, reservation string) (EnvironmentInputReservation, error) + // CreateInputReservation reserves batch under key until the deadline the + // database clock sets, and returns the pending reservation as stored. + CreateInputReservation(ctx context.Context, key string, batch json.RawMessage) (EnvironmentInputReservation, error) + // ExpireInputReservation expires the reservation while it is pending and + // its deadline has passed by the database clock; otherwise it changes + // nothing. + ExpireInputReservation(ctx context.Context, reservation string) error + // AdmitInputReservation settles the pending reservation as admitted and + // returns when the database clock settled it. + AdmitInputReservation(ctx context.Context, reservation string) (time.Time, error) + // FailInputReservation settles the reservation as failed with code while + // it is pending; otherwise it changes nothing. + FailInputReservation(ctx context.Context, reservation, code string) error + // RecordInputAudit records the send_events write audit of the Session. + RecordInputAudit(ctx context.Context) error +} + +// InputStorage is the pooled storage of the input use cases. +type InputStorage interface { + // WithInputs runs apply in one pooled transaction under the lock of the + // tenant's Session and commits only when apply succeeds. A malformed + // tenant is ErrInvalidInput; a malformed, missing or deleted Session is + // ErrNotFound. + WithInputs(ctx context.Context, tenant, session string, apply func(context.Context, InputTx) error) error +} + +// InputReader reads a Session's admitted inputs, its Environment input +// reservations and the reservations execution may admit. +type InputReader interface { + // ListTurnInputs pages a Turn's admitted inputs in sequence order after + // sequence after, at most limit of 1..100. It is an internal recovery + // read, not the public event stream. A malformed ID, a negative after and + // a limit out of range are ErrInvalidInput; a missing Turn is ErrNotFound. + ListTurnInputs(ctx context.Context, tenant, session, turn string, after int64, limit int) ([]TurnInput, error) + // GetEnvironmentInputReservation reads one of the tenant's Session's + // Environment input reservations; an admitted one carries the receipts of + // its batch. A malformed tenant or reservation ID is ErrInvalidInput; a + // missing or deleted Session and a missing reservation are ErrNotFound. + GetEnvironmentInputReservation(ctx context.Context, tenant, session, reservation string) (EnvironmentInputReservation, error) + // ListEnvironmentInputWork lists, in ID order after the reservation ID + // after or from the first when it is empty, at most 100 pending + // reservations before their deadline of live Sessions whose tenant has an + // unrevoked device among connectedDevices, the Session's bound device when + // it has one. A malformed ID is ErrInvalidInput. + ListEnvironmentInputWork(ctx context.Context, after string, connectedDevices []string) ([]EnvironmentInputWork, error) +} + +// SubmitInputs admits a request's input batch in order under one Session +// lock and returns a receipt per input. The whole batch under its key is the +// retry identity: a retry replays the receipts of the batch it admitted, +// whatever has changed since, and the same key with a different batch is +// ErrIdempotencyConflict. A batch with a message waits for a pending file +// write to the Session's Environment (ErrTurnConflict), and no batch is +// admitted while an Environment input reservation is pending or holds the key +// (ErrInputPending). Internal receipts are not the response body of the +// public events endpoint. +func (s *Service) SubmitInputs(ctx context.Context, tenant, session, key string, inputs []Input) ([]InputReceipt, error) { + if err := ValidateInputKey(key); err != nil { + return nil, err + } + batch, encoded, err := ValidateInputs(inputs) + if err != nil { + return nil, err + } + var receipts []InputReceipt + err = s.storage.WithInputs(ctx, tenant, session, func(ctx context.Context, tx InputTx) error { + previous, matches, err := tx.LoadInputBatch(ctx, key, encoded) + if err != nil { + return err + } + replay, err := decideReplay(len(previous) > 0, matches) + if err != nil { + return err + } + if replay { + receipts = previous + return tx.RecordInputAudit(ctx) + } + if slices.ContainsFunc(batch, func(input Input) bool { return input.Kind == "message" }) { + if err := CheckFileWriteGate(ctx, tx); err != nil { + return err + } + } + gate, err := tx.LoadInputGate(ctx, key, encoded) + if err != nil { + return err + } + if err := checkInputGate(gate); err != nil { + return err + } + if receipts, err = AdmitInputs(ctx, tx, key, batch); err != nil { + return err + } + return tx.RecordInputAudit(ctx) + }) + if err != nil { + return nil, fmt.Errorf("submit turn inputs: %w", err) + } + return receipts, nil +} + // FunctionCall retains public identity and its opaque execution-adapter reference. type FunctionCall struct { CallID, ExecutorCallID, Name string diff --git a/services/core/internal/sessions/inputs_test.go b/services/core/internal/sessions/inputs_test.go new file mode 100644 index 000000000..6120dfb46 --- /dev/null +++ b/services/core/internal/sessions/inputs_test.go @@ -0,0 +1,322 @@ +package sessions + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "reflect" + "strings" + "testing" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/items" +) + +// fakeInputTx is fakeFunctionTx with the input methods, equally strict: it +// shares fakeTx's call log, and a method whose func is unset fails the test. +type fakeInputTx struct { + *fakeFunctionTx + + createTurn func() (Turn, error) + createTurnInput func() (int64, error) + loadInputBatch func() ([]InputReceipt, bool, error) + loadInputGate func() (InputGate, error) + findInputReservation func() (*EnvironmentInputReservation, bool, error) + loadInputReservation func() (EnvironmentInputReservation, error) + createInputReservation func() (EnvironmentInputReservation, error) + expireInputReservation func() error + admitInputReservation func() (time.Time, error) + failInputReservation func() error + recordInputAudit func() error +} + +var _ InputTx = (*fakeInputTx)(nil) + +func newInputTx(t *testing.T) *fakeInputTx { + return &fakeInputTx{fakeFunctionTx: &fakeFunctionTx{fakeTx: &fakeTx{t: t}}} +} + +func (f *fakeInputTx) CreateTurn(context.Context) (Turn, error) { + f.record("CreateTurn", f.createTurn != nil) + return f.createTurn() +} + +func (f *fakeInputTx) CreateTurnInput(_ context.Context, turn, key string, position int32, input Input) (int64, error) { + f.record("CreateTurnInput", f.createTurnInput != nil, turn, key, fmt.Sprint(position), input.Kind) + return f.createTurnInput() +} + +func (f *fakeInputTx) LoadInputBatch(_ context.Context, key string, batch json.RawMessage) ([]InputReceipt, bool, error) { + f.record("LoadInputBatch", f.loadInputBatch != nil, key, string(batch)) + return f.loadInputBatch() +} + +func (f *fakeInputTx) LoadInputGate(_ context.Context, key string, batch json.RawMessage) (InputGate, error) { + f.record("LoadInputGate", f.loadInputGate != nil, key, string(batch)) + return f.loadInputGate() +} + +func (f *fakeInputTx) FindInputReservation(_ context.Context, key string, batch json.RawMessage) (*EnvironmentInputReservation, bool, error) { + f.record("FindInputReservation", f.findInputReservation != nil, key, string(batch)) + return f.findInputReservation() +} + +func (f *fakeInputTx) LoadInputReservation(_ context.Context, reservation string) (EnvironmentInputReservation, error) { + f.record("LoadInputReservation", f.loadInputReservation != nil, reservation) + return f.loadInputReservation() +} + +func (f *fakeInputTx) CreateInputReservation(_ context.Context, key string, batch json.RawMessage) (EnvironmentInputReservation, error) { + f.record("CreateInputReservation", f.createInputReservation != nil, key, string(batch)) + return f.createInputReservation() +} + +func (f *fakeInputTx) ExpireInputReservation(_ context.Context, reservation string) error { + f.record("ExpireInputReservation", f.expireInputReservation != nil, reservation) + return f.expireInputReservation() +} + +func (f *fakeInputTx) AdmitInputReservation(_ context.Context, reservation string) (time.Time, error) { + f.record("AdmitInputReservation", f.admitInputReservation != nil, reservation) + return f.admitInputReservation() +} + +func (f *fakeInputTx) FailInputReservation(_ context.Context, reservation, code string) error { + f.record("FailInputReservation", f.failInputReservation != nil, reservation, code) + return f.failInputReservation() +} + +func (f *fakeInputTx) RecordInputAudit(context.Context) error { + f.record("RecordInputAudit", f.recordInputAudit != nil) + return f.recordInputAudit() +} + +func (s *fakeStorage) WithInputs(ctx context.Context, tenant, session string, apply func(context.Context, InputTx) error) error { + s.record("WithInputs", s.inputTx != nil, tenant, session) + return apply(ctx, s.inputTx) +} + +// inputKey is the idempotency key of the input batches under test. +const inputKey = "request" + +func messageInput(text string) Input { + payload, _ := json.Marshal(map[string]string{"text": text}) + return Input{Kind: "message", Payload: payload} +} + +var cancelInput = Input{Kind: "cancel", Payload: json.RawMessage(`{}`)} + +// sequences is a fake CreateTurnInput that allocates sequences from first. +func sequences(first int64) func() (int64, error) { + return func() (int64, error) { + first++ + return first - 1, nil + } +} + +// unprojected is a fake LoadInputSource: a message without text, which +// projects no Item. +var unprojected = returns(Source{Turn: testTurn, Kind: "message", Payload: json.RawMessage(`{}`)}) + +func TestValidateInputs(t *testing.T) { + batch, encoded, err := ValidateInputs([]Input{{Kind: "message", Payload: json.RawMessage(` {"text":"hi","a":1} `)}, {Kind: "cancel", Payload: json.RawMessage(` { } `)}}) + if err != nil || string(batch[0].Payload) != `{"a":1,"text":"hi"}` || string(encoded) != `[{"kind":"message","payload":{"a":1,"text":"hi"}},{"kind":"cancel","payload":{}}]` { + t.Fatalf("batch %s, %v", encoded, err) + } + for name, inputs := range map[string][]Input{ + "no inputs": nil, + "65 inputs": make([]Input, 65), + "unsupported kind": {messageInput("ok"), {Kind: "unsupported", Payload: json.RawMessage(`{}`)}}, + "missing payload": {{Kind: "message"}}, + "array payload": {{Kind: "message", Payload: json.RawMessage(`[]`)}}, + "cancel with a target": {{Kind: "cancel", Payload: json.RawMessage(`{"target":"other"}`)}}, + "over 512 KiB": {messageInput(strings.Repeat("x", 300*1024)), messageInput(strings.Repeat("y", 300*1024))}, + } { + if _, _, err := ValidateInputs(inputs); !errors.Is(err, ErrInvalidInput) { + t.Errorf("%s: %v", name, err) + } + } +} + +func TestValidateMessageInputs(t *testing.T) { + if _, encoded, err := ValidateMessageInputs([]Input{messageInput("hi")}); err != nil || string(encoded) != `[{"kind":"message","payload":{"text":"hi"}}]` { + t.Fatalf("batch %s, %v", encoded, err) + } + for name, inputs := range map[string][]Input{"cancel": {messageInput("hi"), cancelInput}, "no inputs": nil} { + if _, _, err := ValidateMessageInputs(inputs); !errors.Is(err, ErrInvalidInput) { + t.Errorf("%s: %v", name, err) + } + } +} + +func TestDecideReplay(t *testing.T) { + for _, test := range []struct { + used, matches, replay bool + err error + }{ + {false, false, false, nil}, + {true, true, true, nil}, + {true, false, false, ErrIdempotencyConflict}, + } { + if replay, err := decideReplay(test.used, test.matches); replay != test.replay || err != test.err { + t.Errorf("used %v matches %v: %v, %v", test.used, test.matches, replay, err) + } + } +} + +func TestCheckInputGate(t *testing.T) { + for gate, want := range map[InputGate]error{ + {Matches: true}: nil, + {Matches: true, Blocked: true}: ErrInputPending, + {Matches: false, Blocked: true}: ErrIdempotencyConflict, + } { + if err := checkInputGate(gate); err != want { + t.Errorf("%+v: %v", gate, err) + } + } +} + +func TestPlaceInput(t *testing.T) { + for _, test := range []struct { + kind string + active bool + want inputPlacement + }{ + {"message", true, steersTurn}, + {"cancel", true, cancelsTurn}, + {"message", false, startsTurn}, + {"cancel", false, keepsIdentity}, + } { + if got := placeInput(test.kind, test.active); got != test.want { + t.Errorf("%s active %v: %v", test.kind, test.active, got) + } + } +} + +func TestAdmitInput(t *testing.T) { + running := turnWith(TurnInProgress) + t.Run("an idle message starts a Turn", func(t *testing.T) { + var changes []SessionChange + tx := newInputTx(t) + tx.loadActiveTurn, tx.createTurn, tx.appendChanges = activeTurn(nil), returns(turnWith(TurnQueued)), collect(&changes) + tx.createTurnInput, tx.loadUsage = sequences(7), returns(json.RawMessage(`{}`)) + tx.loadInputSource = returns(Source{Turn: testTurn, Kind: "message", Sequence: 7, Payload: json.RawMessage(`{"text":"hi"}`)}) + tx.loadItem, tx.putItem = returns(items.Stored{}), func(items.Change) (*int32, error) { return nil, nil } + receipt, err := admitInput(t.Context(), tx, inputKey, 0, messageInput("hi")) + if err != nil || receipt != (InputReceipt{Sequence: 7, TurnID: testTurn}) { + t.Fatalf("receipt %+v, %v", receipt, err) + } + item := items.Identity(testTurn, "input:7") + assertCalls(t, tx.fakeTx, "LoadActiveTurn", "CreateTurn", "AppendChanges agent.session.turn.created", "CreateTurnInput "+testTurn+" request 0 message", + "LoadInputSource 7", "LoadItem "+testTurn+" "+item, "PutItem "+testTurn+" "+item, "AppendChanges agent.session.turn.item.added", + "LoadUsage", "AppendChanges agent.session.in_progress") + }) + for _, test := range []struct { + name string + input Input + active *Turn + want InputReceipt + calls []string + }{ + {"a message steers the active Turn", messageInput("more"), &running, InputReceipt{Sequence: 8, TurnID: testTurn}, + []string{"LoadActiveTurn", "CreateTurnInput " + testTurn + " request 1 message", "LoadInputSource 8"}}, + {"a cancellation requests the active Turn's cancellation", cancelInput, &running, InputReceipt{Sequence: 8, TurnID: testTurn}, + []string{"LoadActiveTurn", "CreateTurnInput " + testTurn + " request 1 cancel", "RequestTurnCancel " + testTurn, "LoadInputSource 8"}}, + {"an idle cancellation keeps only its retry identity", cancelInput, nil, InputReceipt{Sequence: 8}, + []string{"LoadActiveTurn", "CreateTurnInput request 1 cancel", "LoadInputSource 8"}}, + } { + t.Run(test.name, func(t *testing.T) { + tx := newInputTx(t) + tx.loadActiveTurn, tx.createTurnInput, tx.loadInputSource, tx.requestTurnCancel = activeTurn(test.active), sequences(8), unprojected, done + receipt, err := admitInput(t.Context(), tx, inputKey, 1, test.input) + if err != nil || receipt != test.want { + t.Fatalf("receipt %+v, %v", receipt, err) + } + assertCalls(t, tx.fakeTx, test.calls...) + }) + } + t.Run("a function result joins the Turn of its call", func(t *testing.T) { + tx := newInputTx(t) + tx.findResultTurn = func() (Turn, bool, error) { return turnWith(TurnWaiting), true, nil } + tx.matchFunctionResult, tx.submitFunctionResult, tx.createTurnInput = returns(FunctionResultMatch{Recorded: true}), done, sequences(9) + payload := json.RawMessage(`{"turn_id":"` + testTurn + `","call_id":"call","result":{"ok":true}}`) + receipt, err := admitInput(t.Context(), tx, inputKey, 0, Input{Kind: "tool_result", Payload: payload}) + if err != nil || receipt != (InputReceipt{Sequence: 9, TurnID: testTurn}) { + t.Fatalf("receipt %+v, %v", receipt, err) + } + assertCalls(t, tx.fakeTx, "FindResultTurn "+testTurn, "MatchFunctionResult "+testTurn+` call {"ok":true}`, "SubmitFunctionResult "+testTurn+` call {"ok":true}`, + "CreateTurnInput "+testTurn+" request 0 tool_result") + }) +} + +func TestSubmitInputs(t *testing.T) { + const batch = `[{"kind":"message","payload":{"text":"hi"}},{"kind":"cancel","payload":{}}]` + inputs := []Input{messageInput("hi"), cancelInput} + running := turnWith(TurnInProgress) + submit := func(t *testing.T, tx *fakeInputTx, inputs []Input) ([]InputReceipt, error) { + storage := &fakeStorage{t: t, inputTx: tx} + receipts, err := deviceService(t, storage).SubmitInputs(t.Context(), testTenant, testSession, inputKey, inputs) + assertStorageCalls(t, storage, "WithInputs "+testTenant+" "+testSession) + return receipts, err + } + t.Run("admits the batch in order", func(t *testing.T) { + tx := newInputTx(t) + tx.loadInputBatch = func() ([]InputReceipt, bool, error) { return nil, true, nil } + tx.loadPendingFileWrite, tx.loadInputGate = returns(false), returns(InputGate{Matches: true}) + tx.loadActiveTurn, tx.createTurnInput, tx.loadInputSource, tx.requestTurnCancel = activeTurn(&running), sequences(1), unprojected, done + tx.recordInputAudit = done + receipts, err := submit(t, tx, inputs) + if err != nil || !reflect.DeepEqual(receipts, []InputReceipt{{Sequence: 1, TurnID: testTurn}, {Sequence: 2, TurnID: testTurn}}) { + t.Fatalf("receipts %+v, %v", receipts, err) + } + assertCalls(t, tx.fakeTx, "LoadInputBatch request "+batch, "LoadPendingFileWrite", "LoadInputGate request "+batch, + "LoadActiveTurn", "CreateTurnInput "+testTurn+" request 0 message", "LoadInputSource 1", + "LoadActiveTurn", "CreateTurnInput "+testTurn+" request 1 cancel", "RequestTurnCancel "+testTurn, "LoadInputSource 2", + "RecordInputAudit") + }) + t.Run("a retry replays its receipts", func(t *testing.T) { + previous := []InputReceipt{{Sequence: 1, TurnID: testTurn, Replayed: true}, {Sequence: 2, TurnID: testTurn, Replayed: true}} + tx := newInputTx(t) + tx.loadInputBatch = func() ([]InputReceipt, bool, error) { return previous, true, nil } + tx.recordInputAudit = done + receipts, err := submit(t, tx, inputs) + if err != nil || !reflect.DeepEqual(receipts, previous) { + t.Fatalf("receipts %+v, %v", receipts, err) + } + assertCalls(t, tx.fakeTx, "LoadInputBatch request "+batch, "RecordInputAudit") + }) + t.Run("a different batch under the key conflicts", func(t *testing.T) { + tx := newInputTx(t) + tx.loadInputBatch = func() ([]InputReceipt, bool, error) { return []InputReceipt{{Sequence: 1}}, false, nil } + if _, err := submit(t, tx, inputs); !errors.Is(err, ErrIdempotencyConflict) { + t.Fatal(err) + } + assertCalls(t, tx.fakeTx, "LoadInputBatch request "+batch) + }) + t.Run("a pending file write holds a message", func(t *testing.T) { + tx := newInputTx(t) + tx.loadInputBatch, tx.loadPendingFileWrite = func() ([]InputReceipt, bool, error) { return nil, true, nil }, returns(true) + if _, err := submit(t, tx, inputs); !errors.Is(err, ErrTurnConflict) { + t.Fatal(err) + } + assertCalls(t, tx.fakeTx, "LoadInputBatch request "+batch, "LoadPendingFileWrite") + }) + t.Run("a pending reservation holds a cancellation", func(t *testing.T) { + tx := newInputTx(t) + tx.loadInputBatch, tx.loadInputGate = func() ([]InputReceipt, bool, error) { return nil, true, nil }, returns(InputGate{Matches: true, Blocked: true}) + if _, err := submit(t, tx, []Input{cancelInput}); !errors.Is(err, ErrInputPending) { + t.Fatal(err) + } + assertCalls(t, tx.fakeTx, `LoadInputBatch request [{"kind":"cancel","payload":{}}]`, `LoadInputGate request [{"kind":"cancel","payload":{}}]`) + }) + t.Run("an invalid request touches no storage", func(t *testing.T) { + service := deviceService(t, &fakeStorage{t: t}) + if _, err := service.SubmitInputs(t.Context(), testTenant, testSession, " ", inputs); !errors.Is(err, ErrInvalidInput) { + t.Fatal(err) + } + if _, err := service.SubmitInputs(t.Context(), testTenant, testSession, inputKey, nil); !errors.Is(err, ErrInvalidInput) { + t.Fatal(err) + } + }) +} diff --git a/services/core/internal/sessions/reader.go b/services/core/internal/sessions/reader.go index 6da1497a6..59e24ba32 100644 --- a/services/core/internal/sessions/reader.go +++ b/services/core/internal/sessions/reader.go @@ -9,6 +9,7 @@ type Reader interface { DeviceReader EnvironmentReader ExecutorCredentialReader + InputReader ItemReader ModelExecutionReader SessionReader diff --git a/services/core/internal/sessions/service.go b/services/core/internal/sessions/service.go index 8612a15e7..df891e54c 100644 --- a/services/core/internal/sessions/service.go +++ b/services/core/internal/sessions/service.go @@ -20,4 +20,5 @@ type Storage interface { DeviceStorage ExecutorCredentialStorage SessionStorage + InputStorage } diff --git a/services/core/internal/store/admin_session_archive_race_test.go b/services/core/internal/store/admin_session_archive_race_test.go index 6903443af..b5670acc4 100644 --- a/services/core/internal/store/admin_session_archive_race_test.go +++ b/services/core/internal/store/admin_session_archive_race_test.go @@ -59,7 +59,7 @@ func TestManagedSessionArchiveOrdersConcurrentInput(t *testing.T) { go func() { defer wg.Done() <-start - _, err := s.ReserveEnvironmentInput(t.Context(), tenant, session.ID, "racing-input", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"racing"}`)}}) + _, err := sessionService(t, s).ReserveEnvironmentInput(t.Context(), tenant, session.ID, "racing-input", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"racing"}`)}}) if err != nil && !errors.Is(err, sessions.ErrEnvironmentUnavailable) { t.Error(err) } diff --git a/services/core/internal/store/admin_session_archive_test.go b/services/core/internal/store/admin_session_archive_test.go index d35a6477a..cda2e569b 100644 --- a/services/core/internal/store/admin_session_archive_test.go +++ b/services/core/internal/store/admin_session_archive_test.go @@ -114,7 +114,7 @@ func TestManagedSessionArchiveUnallocatedAndGuards(t *testing.T) { if _, err := deploymentExecution(t, w).ReserveAllocation(t.Context(), deployment.AllocationKey{TenantID: tenant, EnvironmentID: session.Environment.ID}, installation, runtimedevice.HashCredential(uuid.NewString())); !errors.Is(err, deployment.ErrInvalidInput) { t.Fatal("archived Environment allocated after archive", err) } - if _, err := s.ReserveEnvironmentInput(t.Context(), tenant, session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}); !errors.Is(err, sessions.ErrEnvironmentUnavailable) { + if _, err := sessionService(t, s).ReserveEnvironmentInput(t.Context(), tenant, session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}); !errors.Is(err, sessions.ErrEnvironmentUnavailable) { t.Fatal("archived Environment accepted new input", err) } view, err := deploymentService(t, s).View(t.Context()) diff --git a/services/core/internal/store/admin_session_archive_worker_http_test.go b/services/core/internal/store/admin_session_archive_worker_http_test.go index 3c86b135c..7085b64fa 100644 --- a/services/core/internal/store/admin_session_archive_worker_http_test.go +++ b/services/core/internal/store/admin_session_archive_worker_http_test.go @@ -96,7 +96,7 @@ func TestAdminSessionArchiveWorkerHTTPPostgres(t *testing.T) { if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(t.Context(), project.TenantID, active.ID, "pending-turn", json.RawMessage(`{"text":"pending"}`)) + input, err := store.SendMessage(t.Context(), s, project.TenantID, active.ID, "pending-turn", json.RawMessage(`{"text":"pending"}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/archive_cancellation_test.go b/services/core/internal/store/archive_cancellation_test.go index 9c56f82eb..6e3da3dce 100644 --- a/services/core/internal/store/archive_cancellation_test.go +++ b/services/core/internal/store/archive_cancellation_test.go @@ -123,7 +123,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { } time.Sleep(time.Millisecond) } - pending, err := s.ReserveEnvironmentInput(t.Context(), h.tenant, session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}, {Kind: "message", Payload: json.RawMessage(`{"text":"second"}`)}}) + pending, err := store.SessionService(t, s).ReserveEnvironmentInput(t.Context(), h.tenant, session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}, {Kind: "message", Payload: json.RawMessage(`{"text":"second"}`)}}) if err != nil { t.Fatal(err) } @@ -142,7 +142,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { } if scenario == "cancel_revoke_archive" { - if _, err := s.RequestCancel(t.Context(), h.tenant, session.ID, "ordinary-cancel"); err != nil { + if _, err := store.RequestCancel(t.Context(), s, h.tenant, session.ID, "ordinary-cancel"); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/claude_execution_test.go b/services/core/internal/store/claude_execution_test.go index b76101ae7..cf3b54681 100644 --- a/services/core/internal/store/claude_execution_test.go +++ b/services/core/internal/store/claude_execution_test.go @@ -152,7 +152,7 @@ func TestClaudeInvalidImageResultRejectsWholeBatchBeforePersistence(t *testing.T if err != nil || turn.Status != sessions.TurnWaiting || !turn.CancelRequestedAt.IsZero() { t.Fatal(turn, err) } - history, err := h.s.ListTurnInputs(t.Context(), h.tenant, h.session.ID, input.TurnID, 0, 100) + history, err := store.SessionAdapter(h.s).ListTurnInputs(t.Context(), h.tenant, h.session.ID, input.TurnID, 0, 100) if err != nil || len(history) != 1 { t.Fatal(history, err) } diff --git a/services/core/internal/store/command_output_test.go b/services/core/internal/store/command_output_test.go index e2e1cb94f..a87885fa2 100644 --- a/services/core/internal/store/command_output_test.go +++ b/services/core/internal/store/command_output_test.go @@ -22,7 +22,7 @@ func TestCommandOutputCommitsFragmentsSnapshotsAndRecovery(t *testing.T) { if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(ctx, tenant, session.ID, "start", json.RawMessage(`{"text":"run commands"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "start", json.RawMessage(`{"text":"run commands"}`)) if err != nil { t.Fatal(err) } @@ -135,7 +135,7 @@ func TestExecutionJournalsCommandOutputBeforeCancellation(t *testing.T) { h.read(testExecutionRequest) h.write(input.TurnID, proto.TypeToolCall, proto.ToolCallPayload{ID: "cmd", Stage: "before", Observation: &proto.ToolObservation{Kind: "command", Command: "wait", Status: "in_progress"}}) h.write(input.TurnID, proto.TypeCommandOutput, proto.CommandOutputPayload{ID: "cmd", Delta: "partial"}) - if _, err := h.s.RequestCancel(ctx, h.tenant, h.session.ID, "cancel"); err != nil { + if _, err := store.RequestCancel(ctx, h.s, h.tenant, h.session.ID, "cancel"); err != nil { t.Fatal(err) } env := h.read(proto.TypePromptCancel) diff --git a/services/core/internal/store/creation_stream_settlement_public_test.go b/services/core/internal/store/creation_stream_settlement_public_test.go index 7de1cccaf..2dced38ca 100644 --- a/services/core/internal/store/creation_stream_settlement_public_test.go +++ b/services/core/internal/store/creation_stream_settlement_public_test.go @@ -171,7 +171,7 @@ func TestCreationStreamPublicLifetimes(t *testing.T) { created.ended(t, 5*time.Second) connect(first.Session.Environment.ID) - if _, err := s.ReserveEnvironmentInput(t.Context(), tenant, first.Session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}); err != nil { + if _, err := store.SessionService(t, s).ReserveEnvironmentInput(t.Context(), tenant, first.Session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}); err != nil { t.Fatal(err) } if current, err := store.SessionAdapter(s).GetSession(t.Context(), tenant, first.Session.ID); err != nil || !current.PendingInput { @@ -211,7 +211,7 @@ func TestCreationStreamPublicLifetimes(t *testing.T) { if err != nil { t.Fatal(err) } - if settled, err := s.CancelEnvironmentInput(t.Context(), tenant, session, reservation); err != nil || settled.State != sessions.EnvironmentInputCancelled { + if settled, err := store.CancelEnvironmentInput(t.Context(), s, tenant, session, reservation); err != nil || settled.State != sessions.EnvironmentInputCancelled { t.Fatal(settled.State, err) } if after, err := store.SessionAdapter(s).SessionEventCursor(t.Context(), tenant, session); err != nil || after != cursor { diff --git a/services/core/internal/store/deployment_model_providers_http_test.go b/services/core/internal/store/deployment_model_providers_http_test.go index 0b0848aa9..774319df1 100644 --- a/services/core/internal/store/deployment_model_providers_http_test.go +++ b/services/core/internal/store/deployment_model_providers_http_test.go @@ -244,7 +244,7 @@ func TestLegacySessionWithoutProviderCannotStartWork(t *testing.T) { } executor := connectFixtureRuntime(t, h, legacy) // Reserved directly, as a pre-upgrade Core did. - pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, legacy.ID, "before-upgrade", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"old"}`)}}) + pending, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, legacy.ID, "before-upgrade", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"old"}`)}}) if err != nil { t.Fatal(err) } @@ -268,7 +268,7 @@ func TestLegacySessionWithoutProviderCannotStartWork(t *testing.T) { t.Fatal("rejected work was queued", before, after) } awaitDaemonRemoteCondition(t, t.Context(), 5*time.Second, "legacy reservation settled", func() bool { - got, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, legacy.ID, pending.ID) + got, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, legacy.ID, pending.ID) return err == nil && got.State == sessions.EnvironmentInputFailed }) session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, legacy.ID) diff --git a/services/core/internal/store/dispatch_test.go b/services/core/internal/store/dispatch_test.go index e5ca9514c..90d97092a 100644 --- a/services/core/internal/store/dispatch_test.go +++ b/services/core/internal/store/dispatch_test.go @@ -127,7 +127,7 @@ func newDispatchHarnessForSession(t *testing.T, configuration []byte, local bool func (h *dispatchHarness) message(key, text string) sessions.InputReceipt { h.t.Helper() body, _ := json.Marshal(map[string]string{"text": text}) - r, err := h.s.SubmitMessage(context.Background(), h.tenant, h.session.ID, key, body) + r, err := store.SendMessage(context.Background(), h.s, h.tenant, h.session.ID, key, body) if err != nil { h.t.Fatal(err) } @@ -290,7 +290,7 @@ func TestExecutionCancellationRequiresReceiptAndSurvivesContextEnd(t *testing.T) first := h.message("first", "Run") result := h.run(context.Background(), first.TurnID) h.read(testExecutionRequest) - _, err := h.s.RequestCancel(context.Background(), h.tenant, h.session.ID, "cancel") + _, err := store.RequestCancel(context.Background(), h.s, h.tenant, h.session.ID, "cancel") if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_admission_test.go b/services/core/internal/store/environment_admission_test.go index 48f5a5d89..052666fdb 100644 --- a/services/core/internal/store/environment_admission_test.go +++ b/services/core/internal/store/environment_admission_test.go @@ -73,7 +73,7 @@ func environmentAdmissionPending(t *testing.T, h *dispatchHarness, key string) s awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "input reservation", func() bool { return pool.QueryRow(t.Context(), "SELECT id::text FROM environment_input_reservations WHERE session_id=$1 AND idempotency_key=$2", h.session.ID, key).Scan(&id) == nil }) - pending, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, id) + pending, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, id) if err != nil { t.Fatal(err) } @@ -181,7 +181,7 @@ func TestEnvironmentAdmissionSettlementDoesNotCreateTurn(t *testing.T) { } case "cancelled": expected = execution.ErrEnvironmentInputCancelled - if _, err := h.s.CancelEnvironmentInput(t.Context(), h.tenant, h.session.ID, pending.ID); err != nil { + if _, err := store.CancelEnvironmentInput(t.Context(), h.s, h.tenant, h.session.ID, pending.ID); err != nil { t.Fatal(err) } case "deleted": @@ -220,7 +220,7 @@ func TestEnvironmentAdmissionSettlementDoesNotCreateTurn(t *testing.T) { if name == "deleted" { return } - retained, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) + retained, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_claim_worker_test.go b/services/core/internal/store/environment_claim_worker_test.go index a3b57b5b2..62221cd64 100644 --- a/services/core/internal/store/environment_claim_worker_test.go +++ b/services/core/internal/store/environment_claim_worker_test.go @@ -19,8 +19,7 @@ func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { s, db := newTestStoreDB(t) tenant, pending := newEnvironmentExpiryReservation(t, s) owner := executionOwner(t, db, s) - writer := owner.Store - got, err := writer.PromoteEnvironmentInput(t.Context(), tenant, pending.SessionID, pending.ID) + got, err := owner.Sessions.PromoteEnvironmentInput(t.Context(), tenant, pending.SessionID, pending.ID) if err != nil || len(got.Receipts) != 1 || got.Receipts[0].Replayed { t.Fatal(got, err) } @@ -66,7 +65,7 @@ func TestWorkerReconcilesEnvironmentPromotionBeforeStart(t *testing.T) { if err != nil || turns != 1 || inputs != 1 || queued != 0 { t.Fatal("restart duplicated or requeued prepared work", turns, inputs, queued, err) } - successor := executionOwner(t, db, s).Store + successor := executionOwner(t, db, s).Sessions retry, err := successor.PromoteEnvironmentInput(t.Context(), tenant, pending.SessionID, pending.ID) if deleted { if !errors.Is(err, sessions.ErrNotFound) { diff --git a/services/core/internal/store/environment_device_test.go b/services/core/internal/store/environment_device_test.go index 46801cb71..e56cadc9e 100644 --- a/services/core/internal/store/environment_device_test.go +++ b/services/core/internal/store/environment_device_test.go @@ -61,7 +61,7 @@ func TestWorkerEnvironmentSelectsCapableDeviceWithoutMovingBinding(t *testing.T) } stop() for _, value := range []sessions.EnvironmentInputReservation{pending, bound} { - stored, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, value.SessionID, value.ID) + stored, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, value.SessionID, value.ID) if err != nil || stored.State != sessions.EnvironmentInputPending || !stored.Deadline.Equal(value.Deadline) { t.Fatal("device readiness changed pending input", err) } diff --git a/services/core/internal/store/environment_directory_active_test.go b/services/core/internal/store/environment_directory_active_test.go index 161e59146..2f6db1cf1 100644 --- a/services/core/internal/store/environment_directory_active_test.go +++ b/services/core/internal/store/environment_directory_active_test.go @@ -5,12 +5,13 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestEnvironmentDirectoryActiveRunUsesExistingOwner(t *testing.T) { h, w, environment := directoryWorker(t) awaitFixtureCapabilities(t, h, workerEnvironmentCapabilities()) - pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "execute", []sessions.Input{{Kind: "message", Payload: []byte(`{"text":"work"}`)}}) + pending, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "execute", []sessions.Input{{Kind: "message", Payload: []byte(`{"text":"work"}`)}}) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_expiry_worker_test.go b/services/core/internal/store/environment_expiry_worker_test.go index 9ce19dd15..93be26e40 100644 --- a/services/core/internal/store/environment_expiry_worker_test.go +++ b/services/core/internal/store/environment_expiry_worker_test.go @@ -26,7 +26,7 @@ func newEnvironmentExpiryReservation(t *testing.T, s *store.Store) (string, sess if err != nil { t.Fatal(err) } - pending, err := s.ReserveEnvironmentInput(t.Context(), tenant, session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"wait for the environment"}`)}}) + pending, err := store.SessionService(t, s).ReserveEnvironmentInput(t.Context(), tenant, session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"wait for the environment"}`)}}) if err != nil { t.Fatal(err) } @@ -68,7 +68,7 @@ func waitEnvironmentExpiry(t *testing.T, s *store.Store, tenant string, pending t.Helper() deadline := time.Now().Add(5 * time.Second) for { - got, err := s.GetEnvironmentInputReservation(t.Context(), tenant, pending.SessionID, pending.ID) + got, err := store.SessionAdapter(s).GetEnvironmentInputReservation(t.Context(), tenant, pending.SessionID, pending.ID) if err != nil { t.Fatal(err) } @@ -107,7 +107,7 @@ func TestWorkerEnvironmentExpiryWithoutDevicesAndAfterRestart(t *testing.T) { d := &execution.Dispatcher{Store: s, Registry: runtimegateway.NewRegistry()} _, stop := startEnvironmentExpiryWorker(t, db, d) waitEnvironmentExpiry(t, s, dueTenant, due) - got, err := s.GetEnvironmentInputReservation(t.Context(), futureTenant, future.SessionID, future.ID) + got, err := store.SessionAdapter(s).GetEnvironmentInputReservation(t.Context(), futureTenant, future.SessionID, future.ID) if err != nil || got.State != sessions.EnvironmentInputPending || !got.Deadline.Equal(future.Deadline) { t.Fatal("future input changed", got, err) } diff --git a/services/core/internal/store/environment_file_writes_test.go b/services/core/internal/store/environment_file_writes_test.go index 68cc9c7c6..73c342b23 100644 --- a/services/core/internal/store/environment_file_writes_test.go +++ b/services/core/internal/store/environment_file_writes_test.go @@ -61,10 +61,10 @@ func TestEnvironmentFileWriteRetainsUnknownAcrossLeaseLoss(t *testing.T) { if _, err := next.ReserveEnvironmentFileWrite(ctx, f.tenant, f.env.ID, another); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("restart admitted successor", err) } - if _, err := reopened.ReserveEnvironmentInput(ctx, f.tenant, f.session.ID, "new-input", []sessions.Input{messageInput("new")}); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := sessionService(t, reopened).ReserveEnvironmentInput(ctx, f.tenant, f.session.ID, "new-input", []sessions.Input{messageInput("new")}); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("unknown write admitted input", err) } - if _, err := reopened.SubmitMessage(ctx, f.tenant, f.session.ID, "direct", messageInput("new").Payload); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := sendMessage(ctx, reopened, f.tenant, f.session.ID, "direct", messageInput("new").Payload); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("direct admission bypassed write", err) } if _, err := sessionAdapter(reopened).GetEnvironment(ctx, f.tenant, f.env.ID); err != nil { @@ -73,7 +73,7 @@ func TestEnvironmentFileWriteRetainsUnknownAcrossLeaseLoss(t *testing.T) { if _, err := sessionAdapter(reopened).GetSession(ctx, f.tenant, f.session.ID); err != nil { t.Fatal("write gate prevented recovery read", err) } - if receipt, err := reopened.RequestCancel(ctx, f.tenant, f.session.ID, "idle-cancel"); err != nil || receipt.TurnID != "" { + if receipt, err := requestCancel(ctx, reopened, 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 := FixtureFileWrite(ctx, reopened.pool, f.tenant, f.env.ID, f.key.ID); err != nil || got.State != "pending" { @@ -178,7 +178,7 @@ func TestEnvironmentFileWriteSerializesWithInputAndRetry(t *testing.T) { group.Go(func() { <-start var err error - pending, err = f.s.ReserveEnvironmentInput(ctx, f.tenant, f.session.ID, inputKey, []sessions.Input{messageInput("race")}) + pending, err = sessionService(t, f.s).ReserveEnvironmentInput(ctx, f.tenant, f.session.ID, inputKey, []sessions.Input{messageInput("race")}) inputs <- err }) close(start) @@ -189,7 +189,7 @@ func TestEnvironmentFileWriteSerializesWithInputAndRetry(t *testing.T) { t.Fatal(err) } } else if inputErr == nil && errors.Is(writeErr, sessions.ErrTurnConflict) { - if _, err := f.s.CancelEnvironmentInput(ctx, f.tenant, f.session.ID, pending.ID); err != nil { + if _, err := cancelEnvironmentInput(ctx, f.s, f.tenant, f.session.ID, pending.ID); err != nil { t.Fatal(err) } } else { diff --git a/services/core/internal/store/environment_initial_input_test.go b/services/core/internal/store/environment_initial_input_test.go index ae58bba5b..2cb963652 100644 --- a/services/core/internal/store/environment_initial_input_test.go +++ b/services/core/internal/store/environment_initial_input_test.go @@ -21,7 +21,7 @@ func initialEnvironmentReservation(t *testing.T, s *Store, pool *pgxpool.Pool, t if err := pool.QueryRow(t.Context(), "SELECT id FROM environment_input_reservations WHERE session_id=$1 AND is_initial", session).Scan(&id); err != nil { t.Fatal(err) } - reservation, err := s.GetEnvironmentInputReservation(t.Context(), tenant, session, id) + reservation, err := sessionAdapter(s).GetEnvironmentInputReservation(t.Context(), tenant, session, id) if err != nil || !reservation.IsInitial { t.Fatal("missing initial origin", reservation, err) } @@ -49,7 +49,7 @@ func TestEnvironmentInitialExpiryRollsBackWithFailureEventAndSerializesPromotion _, _ = pool.Exec(context.Background(), "ALTER TABLE session_events DROP CONSTRAINT IF EXISTS "+constraint) }) writer := executionWriter(t, s) - if _, err := writer.ExpireEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID); err == nil { + if _, err := sessionService(t, writer).ExpireEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID); err == nil { t.Fatal("expiry committed without its failure event") } retained := initialEnvironmentReservation(t, s, pool, tenant, session.ID) @@ -65,7 +65,7 @@ func TestEnvironmentInitialExpiryRollsBackWithFailureEventAndSerializesPromotion err error } results := make(chan result, 2) - for _, settle := range []func(context.Context, string, string, string) (sessions.EnvironmentInputReservation, error){writer.PromoteEnvironmentInput, s.ExpireEnvironmentInput} { + for _, settle := range []func(context.Context, string, string, string) (sessions.EnvironmentInputReservation, error){sessionExecution(t, writer.lease).PromoteEnvironmentInput, sessionService(t, s).ExpireEnvironmentInput} { go func() { r, err := settle(t.Context(), tenant, session.ID, reservation.ID); results <- result{r, err} }() } for i := 0; i < 2; i++ { @@ -156,7 +156,7 @@ func TestEnvironmentInitialInputCreationRetainsCursorIdentityAndPromotion(t *tes } requireEnvironmentInputActivity(t, s, tenant, session.ID, connectedStatus, "") environmentInputHistory(t, pool, session.ID, 0, 0) - promoted, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) + promoted, err := sessionExecution(t, writer.lease).PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) if err != nil || promoted.State != sessions.EnvironmentInputAdmitted || !promoted.IsInitial || len(promoted.Receipts) != 2 { t.Fatal("initial batch did not promote", promoted, err) } @@ -164,7 +164,7 @@ func TestEnvironmentInitialInputCreationRetainsCursorIdentityAndPromotion(t *tes if active.LastTurn == nil || active.LastTurn.Status != sessions.TurnInProgress || active.PendingInput { t.Fatal("promotion did not claim its Turn", active.LastTurn) } - replay, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) + replay, err := sessionExecution(t, writer.lease).PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) if err != nil || len(replay.Receipts) != 2 || !replay.Receipts[0].Replayed || !replay.Receipts[1].Replayed { t.Fatal("promotion retry granted fresh receipts", replay, err) } @@ -211,7 +211,7 @@ func TestEnvironmentInitialInputExpiryHasNoTurnAndCannotReplay(t *testing.T) { } writer := executionWriter(t, s) for reservation.State == sessions.EnvironmentInputPending { - count, err := writer.ExpireEnvironmentInputs(t.Context()) + count, err := sessionExecution(t, writer.lease).ExpireEnvironmentInputs(t.Context()) if err != nil || count < 1 || count > 32 { t.Fatal("expiry made no bounded progress", count, err) } @@ -250,7 +250,7 @@ func TestEnvironmentInitialInputExpiryHasNoTurnAndCannotReplay(t *testing.T) { if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, session.Environment.ID, generation, 1, true); err != nil { t.Fatal(err) } - late, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) + late, err := sessionExecution(t, writer.lease).PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) if err != nil || late.State != sessions.EnvironmentInputExpired || len(late.Receipts) != 0 { t.Fatal("late connection resurrected initial input", late, err) } @@ -269,7 +269,7 @@ func TestEnvironmentInitialInputExpiryHasNoTurnAndCannotReplay(t *testing.T) { if err := sessionService(t, reopened).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); !errors.Is(err, sessions.ErrNotIdle) { t.Fatal("pending later input deleted", err) } - if _, err := reopened.CancelEnvironmentInput(t.Context(), tenant, session.ID, later.ID); err != nil { + if _, err := cancelEnvironmentInput(t.Context(), reopened, tenant, session.ID, later.ID); err != nil { t.Fatal(err) } if err := sessionService(t, reopened).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { diff --git a/services/core/internal/store/environment_initial_public_test.go b/services/core/internal/store/environment_initial_public_test.go index 82deb5708..dd490c24a 100644 --- a/services/core/internal/store/environment_initial_public_test.go +++ b/services/core/internal/store/environment_initial_public_test.go @@ -38,7 +38,6 @@ func TestEnvironmentInitialFailureOfficialClient(t *testing.T) { t.Fatal(err) } owner := executionOwner(t, db, s) - writer := owner.Store t.Cleanup(func() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() @@ -94,7 +93,7 @@ func TestEnvironmentInitialFailureOfficialClient(t *testing.T) { if err := db.pool.QueryRow(t.Context(), "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1 AND is_initial RETURNING id", session.ID).Scan(&reservation); err != nil { t.Fatal(err) } - if result, err := writer.ExpireEnvironmentInput(t.Context(), tenant, session.ID, reservation); err != nil || result.State != sessions.EnvironmentInputExpired { + if result, err := store.SessionService(t, s).ExpireEnvironmentInput(t.Context(), tenant, session.ID, reservation); err != nil || result.State != sessions.EnvironmentInputExpired { t.Fatal("initial reservation did not expire", result, err) } observed := <-done diff --git a/services/core/internal/store/environment_input_activity.go b/services/core/internal/store/environment_input_activity.go deleted file mode 100644 index 5e32b7e0e..000000000 --- a/services/core/internal/store/environment_input_activity.go +++ /dev/null @@ -1,23 +0,0 @@ -package store - -import ( - "context" - - "github.com/jackc/pgx/v5/pgtype" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -// withEnvironmentInputSession runs apply in the public Session's transaction -// and reports the input activity it changes. -func (s *Store) withEnvironmentInputSession(ctx context.Context, tenantID, sessionID string, apply func(context.Context, *sqlc.Queries, pgtype.UUID) error) error { - tenant, err := parseID(tenantID) - if err != nil { - return err - } - return s.withPublicSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, id pgtype.UUID) error { - return sessions.TrackInputActivity(ctx, sessionpg.BindSession(q, tenant, id), func(ctx context.Context) error { return apply(ctx, q, id) }) - }) -} diff --git a/services/core/internal/store/environment_input_activity_test.go b/services/core/internal/store/environment_input_activity_test.go index 2dd6a8cd9..b874348f9 100644 --- a/services/core/internal/store/environment_input_activity_test.go +++ b/services/core/internal/store/environment_input_activity_test.go @@ -85,7 +85,7 @@ func TestEnvironmentInputActivityWaitsBeforeTurnAndClearsOnConnection(t *testing if err := sessionExecution(t, writer.lease).ObserveEnvironmentConnection(t.Context(), tenant, environment, generation, 3, true); err != nil { t.Fatal(err) } - if _, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID); err != nil { + if _, err := sessionExecution(t, writer.lease).PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID); err != nil { t.Fatal(err) } active := requireEnvironmentInputActivity(t, s, tenant, session.ID, "", "") @@ -103,7 +103,7 @@ func TestEnvironmentInputActivitySettlementAndNewerWork(t *testing.T) { s, pool := testStore(t) tenant, session := environmentInputSession(t, s) writer := executionWriter(t, s) - prior, err := s.SubmitInputs(t.Context(), tenant, session.ID, "prior", []sessions.Input{messageInput("prior")}) + prior, err := submitInputs(t.Context(), s, tenant, session.ID, "prior", []sessions.Input{messageInput("prior")}) if err != nil { t.Fatal(err) } @@ -120,18 +120,18 @@ func TestEnvironmentInputActivitySettlementAndNewerWork(t *testing.T) { t.Fatal(err) } if state == sessions.EnvironmentInputCancelled { - _, err = s.CancelEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) + _, err = cancelEnvironmentInput(t.Context(), s, tenant, session.ID, reservation.ID) } else { if _, err := pool.Exec(t.Context(), "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE id=$1", reservation.ID); err != nil { t.Fatal(err) } // Other retained test rows may precede this reservation in bounded batches. for reservation.State == sessions.EnvironmentInputPending { - count, sweepErr := writer.ExpireEnvironmentInputs(t.Context()) + count, sweepErr := sessionExecution(t, writer.lease).ExpireEnvironmentInputs(t.Context()) if sweepErr != nil || count < 1 || count > 32 { t.Fatal("expiry made no bounded progress", count, sweepErr) } - reservation, err = s.GetEnvironmentInputReservation(t.Context(), tenant, session.ID, reservation.ID) + reservation, err = sessionAdapter(s).GetEnvironmentInputReservation(t.Context(), tenant, session.ID, reservation.ID) if err != nil { t.Fatal(err) } @@ -155,7 +155,7 @@ func TestEnvironmentInputActivitySettlementAndNewerWork(t *testing.T) { if after, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID); err != nil || after != cursor { t.Fatal("settled retry repeated activity", after, cursor, err) } - if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "newer", []sessions.Input{messageInput("newer")}); err != nil { + if _, err := submitInputs(t.Context(), s, tenant, session.ID, "newer", []sessions.Input{messageInput("newer")}); err != nil { t.Fatal(err) } requireEnvironmentInputActivity(t, s, tenant, session.ID, "", "") @@ -173,7 +173,7 @@ func TestEnvironmentInputActivityRollsBackReservationAndConnection(t *testing.T) t.Cleanup(func() { _, _ = pool.Exec(context.Background(), "ALTER TABLE session_events DROP CONSTRAINT IF EXISTS "+constraint) }) - if _, err := s.ReserveEnvironmentInput(t.Context(), tenant, session.ID, "rollback", []sessions.Input{messageInput("pending")}); err == nil { + if _, err := sessionService(t, s).ReserveEnvironmentInput(t.Context(), tenant, session.ID, "rollback", []sessions.Input{messageInput("pending")}); err == nil { t.Fatal("activity failure retained reservation") } var count int @@ -227,7 +227,7 @@ func TestEnvironmentInputActivityRecoversWaitingActionAndHidesDeletion(t *testin t.Fatal(err) } requireEnvironmentInputActivity(t, s, tenant, session.ID, "requires_action", environment) - got, err := s.GetEnvironmentInputReservation(t.Context(), tenant, session.ID, reservation.ID) + got, err := sessionAdapter(s).GetEnvironmentInputReservation(t.Context(), tenant, session.ID, reservation.ID) if err != nil || got.State != sessions.EnvironmentInputPending || !got.Deadline.Equal(reservation.Deadline) { t.Fatal("recovery changed waiting input or its deadline", got, err) } @@ -255,25 +255,3 @@ func TestEnvironmentInputActivityRecoversWaitingActionAndHidesDeletion(t *testin t.Fatal("deleted activity events remained visible", err) } } - -func TestPreparationFailurePreservesCancelledAndNewerInput(t *testing.T) { - s, _ := testStore(t) - tenant, session := environmentInputSession(t, s) - first := reserveEnvironmentInput(t, s, tenant, session.ID, "cancelled") - if _, err := s.CancelEnvironmentInput(t.Context(), tenant, session.ID, first.ID); err != nil { - t.Fatal(err) - } - next := reserveEnvironmentInput(t, s, tenant, session.ID, "new") - if err := s.FailEnvironmentInput(t.Context(), tenant, session.ID, first.ID, "runtime_preparation_failed"); err != nil { - t.Fatal(err) - } - for id, state := range map[string]string{first.ID: sessions.EnvironmentInputCancelled, next.ID: sessions.EnvironmentInputPending} { - current, err := s.GetEnvironmentInputReservation(t.Context(), tenant, session.ID, id) - if err != nil || current.State != state { - t.Fatal("late failure changed another outcome", err) - } - } - if err := s.FailEnvironmentInput(t.Context(), tenant, session.ID, next.ID, "secret-canary"); !errors.Is(err, sessions.ErrInvalidInput) { - t.Fatal("unclassified diagnostic accepted", err) - } -} diff --git a/services/core/internal/store/environment_input_claim_test.go b/services/core/internal/store/environment_input_claim_test.go deleted file mode 100644 index 8ce51ee8d..000000000 --- a/services/core/internal/store/environment_input_claim_test.go +++ /dev/null @@ -1,96 +0,0 @@ -package store - -import ( - "sync" - "testing" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -func TestEnvironmentInputConcurrentPromotionClaimsOnce(t *testing.T) { - s, pool := testStore(t) - writer := executionWriter(t, s) - tenant, session := environmentInputSession(t, s) - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - reservationCursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) - if err != nil { - t.Fatal(err) - } - const count = 8 - results := make(chan sessions.EnvironmentInputReservation, count) - var group sync.WaitGroup - for range count { - group.Go(func() { - got, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID) - if err != nil { - t.Error(err) - return - } - results <- got - }) - } - group.Wait() - close(results) - var turnID string - fresh, received := 0, 0 - for got := range results { - received++ - if got.State != sessions.EnvironmentInputAdmitted || len(got.Receipts) != 2 { - t.Fatal("promotion lost the original batch", got) - } - if turnID == "" { - turnID = got.Receipts[0].TurnID - } - if !got.Receipts[0].Replayed { - fresh++ - } - for _, receipt := range got.Receipts { - if receipt.TurnID != turnID || receipt.Replayed != got.Receipts[0].Replayed { - t.Fatal("promotion changed execution ownership", got.Receipts) - } - } - } - if received != count || fresh != 1 { - t.Fatal("promotion authorized multiple starts", received, fresh) - } - turn, err := sessionAdapter(s).GetTurn(t.Context(), tenant, session.ID, turnID) - if err != nil || turn.Status != sessions.TurnInProgress || turn.StartedAt.IsZero() { - t.Fatal("promotion did not persist its execution claim", turn, err) - } - environmentInputHistory(t, pool, session.ID, 1, 2) - changes, err := sessionAdapter(s).ListSessionEvents(t.Context(), tenant, session.ID, reservationCursor) - if err != nil || len(changes) < 2 { - t.Fatal("missing promotion events", changes, err) - } - created, claimed := changes[0], changes[len(changes)-1] - if created.Event.Type != "agent.session.turn.created" || created.Turn == nil || created.Turn.Status != sessions.TurnQueued || claimed.Event.Type != "agent.session.turn.in_progress" || claimed.Turn == nil || claimed.Turn.ID != turnID || claimed.Turn.Status != sessions.TurnInProgress { - t.Fatal("claim reordered or replaced admission snapshots", created, claimed) - } - claimEvents := 0 - for _, change := range changes { - if change.Event.Type == "agent.session.turn.in_progress" { - claimEvents++ - } - } - if claimEvents != 1 { - t.Fatal("retry published another claim", claimEvents) - } - transition(t, writer, tenant, session.ID, turnID, sessions.TurnInProgress, sessions.TurnCompleted) - later := reserveEnvironmentInput(t, s, tenant, session.ID, "later") - cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) - if err != nil { - t.Fatal(err) - } - retry, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID) - if err != nil || len(retry.Receipts) != 2 || !retry.Receipts[0].Replayed || retry.Receipts[0].TurnID != turnID { - t.Fatal("terminal retry reclaimed execution", retry, err) - } - after, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) - if err != nil || after != cursor { - t.Fatal("terminal retry published events", after, cursor, err) - } - retained, err := s.GetEnvironmentInputReservation(t.Context(), tenant, session.ID, later.ID) - if err != nil || retained.State != sessions.EnvironmentInputPending || !retained.Deadline.Equal(later.Deadline) { - t.Fatal("old promotion affected new preparation", retained, err) - } -} diff --git a/services/core/internal/store/environment_input_expiry.go b/services/core/internal/store/environment_input_expiry.go deleted file mode 100644 index 3724937da..000000000 --- a/services/core/internal/store/environment_input_expiry.go +++ /dev/null @@ -1,43 +0,0 @@ -package store - -import ( - "context" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" -) - -// ExpireEnvironmentInputs settles one bounded batch without creating Turn history. -// Only the current execution writer may run this cross-Session maintenance. -func (s *Store) ExpireEnvironmentInputs(ctx context.Context) (int64, error) { - if err := s.checkExecutionAuthority(); err != nil { - return 0, err - } - var expired int64 - err := s.writer.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - rows, err := q.ListDueEnvironmentInputs(ctx) - if err != nil { - return err - } - for _, row := range rows { - err := sessions.TrackInputActivity(ctx, sessionpg.BindSession(q, row.TenantID, row.SessionID), func(ctx context.Context) error { - return q.ExpireEnvironmentInputReservation(ctx, sqlc.ExpireEnvironmentInputReservationParams{SessionID: row.SessionID, ID: row.ID}) - }) - if err != nil { - return err - } - if err := sessionpg.PruneChanges(ctx, q, row.SessionID); err != nil { - return err - } - expired++ - } - return nil - }) - if err != nil { - return 0, err - } - return expired, nil -} diff --git a/services/core/internal/store/environment_input_expiry_test.go b/services/core/internal/store/environment_input_expiry_test.go index f733964f8..4c003e0e6 100644 --- a/services/core/internal/store/environment_input_expiry_test.go +++ b/services/core/internal/store/environment_input_expiry_test.go @@ -10,86 +10,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -func TestEnvironmentExpiryBoundsBatchAndRequiresExecutionWriter(t *testing.T) { - s, pool := testStore(t) - ctx := t.Context() - var reservations []sessions.EnvironmentInputReservation - for range 33 { - tenant, session := environmentInputSession(t, s) - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - reservations = append(reservations, pending) - if _, err := pool.Exec(ctx, "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE id=$1", pending.ID); err != nil { - t.Fatal(err) - } - } - if n, err := s.ExpireEnvironmentInputs(ctx); err == nil || n != 0 { - t.Fatal("unleased maintenance", n, err) - } - var before, after, due int64 - if err := pool.QueryRow(ctx, "SELECT count(*) FILTER (WHERE state='expired'), count(*) FILTER (WHERE state='pending' AND deadline <= statement_timestamp()) FROM environment_input_reservations").Scan(&before, &due); err != nil { - t.Fatal(err) - } - writer := executionWriter(t, s) - n, err := writer.ExpireEnvironmentInputs(ctx) - if err != nil || n != 32 { - t.Fatal("unbounded or incomplete batch", n, err) - } - if err := pool.QueryRow(ctx, "SELECT count(*) FROM environment_input_reservations WHERE state='expired'").Scan(&after); err != nil || after-before != n { - t.Fatal("reported expiry did not match persisted settlement", before, after, n, err) - } - // Existing test rows may precede this fixture; each pass must make bounded progress. - for remaining := due - n; remaining > 0; { - n, err = writer.ExpireEnvironmentInputs(ctx) - if err != nil || n <= 0 || n > 32 { - t.Fatal("expiry backlog did not progress", n, err) - } - remaining -= n - } - for _, pending := range reservations { - var state string - if err := pool.QueryRow(ctx, "SELECT state FROM environment_input_reservations WHERE id=$1", pending.ID).Scan(&state); err != nil || state != sessions.EnvironmentInputExpired { - t.Fatal(state, err) - } - environmentInputHistory(t, pool, pending.SessionID, 0, 0) - } -} - -func TestEnvironmentExpiryFencesLostExecutionOwner(t *testing.T) { - s, pool := testStore(t) - tenant, session := environmentInputSession(t, s) - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - if _, err := pool.Exec(t.Context(), "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE id=$1", pending.ID); err != nil { - t.Fatal(err) - } - old := executionWriter(t, s) - var killed bool - if err := pool.QueryRow(t.Context(), "SELECT pg_terminate_backend($1,1000)", executionOwnerPID(t, pool)).Scan(&killed); err != nil || !killed { - t.Fatal(killed, err) - } - successor := executionWriter(t, s) - if n, err := old.ExpireEnvironmentInputs(t.Context()); err == nil || n != 0 { - t.Fatal("lost owner expired input", n, err) - } - got, err := s.GetEnvironmentInputReservation(t.Context(), tenant, session.ID, pending.ID) - if err != nil || got.State != sessions.EnvironmentInputPending { - t.Fatal("lost owner wrote through the pool", got, err) - } - for got.State == sessions.EnvironmentInputPending { - n, err := successor.ExpireEnvironmentInputs(t.Context()) - if err != nil || n == 0 { - t.Fatal("successor could not expire input", n, err) - } - got, err = s.GetEnvironmentInputReservation(t.Context(), tenant, session.ID, pending.ID) - if err != nil { - t.Fatal(err) - } - } - if got.State != sessions.EnvironmentInputExpired { - t.Fatal(got) - } - environmentInputHistory(t, pool, session.ID, 0, 0) -} - func TestEnvironmentExpirySerializesWithTargetedSettlement(t *testing.T) { for _, action := range []string{"promote", "cancel", "delete"} { t.Run(action, func(t *testing.T) { @@ -99,18 +19,22 @@ func TestEnvironmentExpirySerializesWithTargetedSettlement(t *testing.T) { if _, err := pool.Exec(t.Context(), "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE id=$1", pending.ID); err != nil { t.Fatal(err) } - writer := executionWriter(t, s) + operations := sessionExecution(t, executionWriter(t, s).lease) start := make(chan struct{}) results := make(chan error, 2) - go func() { <-start; _, err := writer.ExpireEnvironmentInputs(t.Context()); results <- err }() + go func() { + <-start + _, err := operations.ExpireEnvironmentInputs(t.Context()) + results <- err + }() go func() { <-start var err error switch action { case "promote": - _, err = writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID) + _, err = operations.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID) case "cancel": - _, err = s.CancelEnvironmentInput(t.Context(), tenant, session.ID, pending.ID) + _, err = cancelEnvironmentInput(t.Context(), s, tenant, session.ID, pending.ID) case "delete": err = sessionService(t, s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}) } @@ -131,22 +55,25 @@ func TestEnvironmentExpirySerializesWithTargetedSettlement(t *testing.T) { } environmentInputHistory(t, pool, session.ID, 0, 0) if action == "delete" { - if _, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := operations.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted input resurrected", err) } return } later := reserveEnvironmentInput(t, s, tenant, session.ID, uuid.NewString()) - for _, settle := range []func(context.Context, string, string, string) (sessions.EnvironmentInputReservation, error){writer.PromoteEnvironmentInput, s.CancelEnvironmentInput, s.ExpireEnvironmentInput} { + for _, settle := range []func(context.Context, string, string, string) (sessions.EnvironmentInputReservation, error){ + operations.PromoteEnvironmentInput, + sessionService(t, s).ExpireEnvironmentInput, + } { old, err := settle(t.Context(), tenant, session.ID, pending.ID) if err != nil || old.State != state { t.Fatal("old reservation changed", old, err) } } - if _, err := writer.ExpireEnvironmentInputs(t.Context()); err != nil { + if _, err := operations.ExpireEnvironmentInputs(t.Context()); err != nil { t.Fatal(err) } - got, err := s.GetEnvironmentInputReservation(t.Context(), tenant, session.ID, later.ID) + got, err := sessionAdapter(s).GetEnvironmentInputReservation(t.Context(), tenant, session.ID, later.ID) if err != nil || got.State != sessions.EnvironmentInputPending || !got.Deadline.Equal(later.Deadline) { t.Fatal("old settlement affected successor", got, err) } diff --git a/services/core/internal/store/environment_input_migration_test.go b/services/core/internal/store/environment_input_migration_test.go index c04de8f7a..a792de72c 100644 --- a/services/core/internal/store/environment_input_migration_test.go +++ b/services/core/internal/store/environment_input_migration_test.go @@ -7,8 +7,6 @@ import ( "strings" "testing" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/stdlib" @@ -92,37 +90,3 @@ func TestEnvironmentInputMigrationRetainsHistoryAndRetryIdentity(t *testing.T) { t.Fatal("migration manufactured a Turn", count, err) } } - -func TestEnvironmentInputPromotionUsesCurrentExecutionWriter(t *testing.T) { - s, pool := testStore(t) - tenant, session := environmentInputSession(t, s) - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - if _, err := s.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID); err == nil { - t.Fatal("pooled Store promoted input without execution ownership") - } - closed := executionWriter(t, s) - awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, closed.pool) - if err := closed.lease.Close(t.Context()); err != nil { - t.Fatal(err) - } - awaitRelease() - if _, err := closed.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID); err == nil { - t.Fatal("closed execution writer promoted pending input") - } - environmentInputHistory(t, pool, session.ID, 0, 0) - writer := executionWriter(t, s) - var killed bool - if err := pool.QueryRow(t.Context(), "SELECT pg_terminate_backend($1, 1000)", executionOwnerPID(t, pool)).Scan(&killed); err != nil || !killed { - t.Fatal(killed, err) - } - successor := executionWriter(t, s) - if _, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID); err == nil { - t.Fatal("stale execution writer promoted pending input") - } - environmentInputHistory(t, pool, session.ID, 0, 0) - got, err := successor.PromoteEnvironmentInput(t.Context(), tenant, session.ID, pending.ID) - if err != nil || got.State != sessions.EnvironmentInputAdmitted { - t.Fatal("successor could not promote", got, err) - } - environmentInputHistory(t, pool, session.ID, 1, 2) -} diff --git a/services/core/internal/store/environment_input_settlement_test.go b/services/core/internal/store/environment_input_settlement_test.go index 874d8958a..ac8ee7b74 100644 --- a/services/core/internal/store/environment_input_settlement_test.go +++ b/services/core/internal/store/environment_input_settlement_test.go @@ -3,223 +3,11 @@ package store import ( "context" "errors" - "strings" "testing" - "time" - - "github.com/google/uuid" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -func TestEnvironmentInputTerminalReservationsCannotRestart(t *testing.T) { - for _, terminal := range []string{sessions.EnvironmentInputCancelled, sessions.EnvironmentInputExpired} { - t.Run(terminal, func(t *testing.T) { - s, pool := testStore(t) - writer := executionWriter(t, s) - tenant, session := environmentInputSession(t, s) - ctx := context.Background() - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - early, err := s.ExpireEnvironmentInput(ctx, tenant, session.ID, pending.ID) - if err != nil || early.State != sessions.EnvironmentInputPending || early.SettledAt != nil || !early.Deadline.Equal(pending.Deadline) { - t.Fatal("early expiry", early, err) - } - var settled sessions.EnvironmentInputReservation - if terminal == sessions.EnvironmentInputExpired { - if _, err := pool.Exec(ctx, "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE id=$1", pending.ID); err != nil { - t.Fatal(err) - } - settled, err = writer.PromoteEnvironmentInput(ctx, tenant, session.ID, pending.ID) - } else { - settled, err = s.CancelEnvironmentInput(ctx, tenant, session.ID, pending.ID) - } - if err != nil || settled.State != terminal || settled.SettledAt == nil || len(settled.Receipts) != 0 { - t.Fatal("terminal settlement", settled, err) - } - retry, err := s.ReserveEnvironmentInput(ctx, tenant, session.ID, "pending", pending.Inputs) - if err != nil || retry.State != terminal || retry.ID != pending.ID || !retry.Deadline.Equal(settled.Deadline) || !retry.SettledAt.Equal(*settled.SettledAt) { - t.Fatal("terminal retry changed outcome", retry, err) - } - if _, err := s.SubmitInputs(ctx, tenant, session.ID, "pending", pending.Inputs); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal("terminal request reopened through direct path", err) - } - if _, err := s.SubmitInputs(ctx, tenant, session.ID, "pending", []sessions.Input{messageInput("changed")}); !errors.Is(err, sessions.ErrIdempotencyConflict) { - t.Fatal("terminal identity changed", err) - } - later := reserveEnvironmentInput(t, s, tenant, session.ID, "later") - for _, finish := range []func(context.Context, string, string, string) (sessions.EnvironmentInputReservation, error){ - writer.PromoteEnvironmentInput, s.CancelEnvironmentInput, s.ExpireEnvironmentInput, - } { - got, err := finish(ctx, tenant, session.ID, pending.ID) - if err != nil || got.State != terminal { - t.Fatal("old settlement changed", got, err) - } - } - got, err := s.GetEnvironmentInputReservation(ctx, tenant, session.ID, later.ID) - if err != nil || got.State != sessions.EnvironmentInputPending || !got.Deadline.Equal(later.Deadline) { - t.Fatal("old settlement touched successor", got, err) - } - environmentInputHistory(t, pool, session.ID, 0, 0) - }) - } -} - -func TestEnvironmentInputPromotionRollsBackHistoryAndSettlement(t *testing.T) { - for _, phase := range []string{"input", "settlement", "claim", "claim-event"} { - t.Run(phase, func(t *testing.T) { - s, pool := testStore(t) - writer := executionWriter(t, s) - tenant, session := environmentInputSession(t, s) - ctx := context.Background() - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - name := "reservation_failure_" + strings.ReplaceAll(uuid.NewString(), "-", "") - table := "turn_inputs" - expression := "session_id <> '" + session.ID + "'::uuid OR payload->>'text' <> 'second'" - if phase == "settlement" { - table = "environment_input_reservations" - expression = "id <> '" + pending.ID + "'::uuid OR state <> 'admitted'" - } - if phase == "claim" { - table = "turns" - expression = "session_id <> '" + session.ID + "'::uuid OR status <> 'in_progress'" - } - if phase == "claim-event" { - table = "session_events" - expression = "session_id <> '" + session.ID + "'::uuid OR payload->'event'->>'type' <> 'agent.session.turn.in_progress'" - } - if _, err := pool.Exec(ctx, "ALTER TABLE "+table+" ADD CONSTRAINT "+name+" CHECK ("+expression+") NOT VALID"); err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _, _ = pool.Exec(ctx, "ALTER TABLE "+table+" DROP CONSTRAINT IF EXISTS "+name) }) - if _, err := writer.PromoteEnvironmentInput(ctx, tenant, session.ID, pending.ID); err == nil { - t.Fatal("injected failure succeeded") - } - environmentInputHistory(t, pool, session.ID, 0, 0) - got, err := s.GetEnvironmentInputReservation(ctx, tenant, session.ID, pending.ID) - if err != nil || got.State != sessions.EnvironmentInputPending || got.SettledAt != nil || !got.Deadline.Equal(pending.Deadline) { - t.Fatal("partial settlement survived", got, err) - } - if _, err := pool.Exec(ctx, "ALTER TABLE "+table+" DROP CONSTRAINT "+name); err != nil { - t.Fatal(err) - } - got, err = writer.PromoteEnvironmentInput(ctx, tenant, session.ID, pending.ID) - if err != nil || got.State != sessions.EnvironmentInputAdmitted { - t.Fatal(got, err) - } - environmentInputHistory(t, pool, session.ID, 1, 2) - }) - } -} - -func TestEnvironmentInputDeadlineIsCheckedAfterSessionLock(t *testing.T) { - for _, action := range []string{"promote", "fail"} { - t.Run(action, func(t *testing.T) { - s, pool := testStore(t) - writer := executionWriter(t, s) - tenant, session := environmentInputSession(t, s) - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - tx, err := pool.Begin(ctx) - if err != nil { - t.Fatal(err) - } - defer func() { _ = tx.Rollback(context.Background()) }() - var blocker int32 - if err := tx.QueryRow(ctx, "SELECT pg_backend_pid() FROM sessions WHERE id=$1 FOR UPDATE", session.ID).Scan(&blocker); err != nil { - t.Fatal(err) - } - type outcome struct { - value sessions.EnvironmentInputReservation - err error - } - done := make(chan outcome, 1) - go func() { - var got sessions.EnvironmentInputReservation - var err error - if action == "promote" { - got, err = writer.PromoteEnvironmentInput(ctx, tenant, session.ID, pending.ID) - } else { - err = writer.FailEnvironmentInput(ctx, tenant, session.ID, pending.ID, "runtime_preparation_failed") - if err == nil { - got, err = s.GetEnvironmentInputReservation(ctx, tenant, session.ID, pending.ID) - } - } - done <- outcome{got, err} - }() - for { - var blocked bool - if err := pool.QueryRow(ctx, "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE $1=ANY(pg_blocking_pids(pid)))", blocker).Scan(&blocked); err != nil { - t.Fatal(err) - } - if blocked { - break - } - select { - case result := <-done: - t.Fatal("settlement bypassed Session lock", result) - case <-ctx.Done(): - t.Fatal("settlement lock wait not observed") - case <-time.After(5 * time.Millisecond): - } - } - // Transaction-start time is now older than the controlled deadline. - if _, err := tx.Exec(ctx, "UPDATE environment_input_reservations SET deadline=clock_timestamp() WHERE id=$1", pending.ID); err != nil { - t.Fatal(err) - } - if err := tx.Commit(ctx); err != nil { - t.Fatal(err) - } - result := <-done - if result.err != nil || result.value.State != sessions.EnvironmentInputExpired || result.value.SettledAt == nil { - t.Fatal("lock wait extended input lifetime", result) - } - environmentInputHistory(t, pool, session.ID, 0, 0) - stored, err := s.GetEnvironmentInputReservation(ctx, tenant, session.ID, pending.ID) - if err != nil || stored.State != sessions.EnvironmentInputExpired { - t.Fatal("expiry was rolled back", stored, err) - } - }) - } -} - -func TestEnvironmentInputCancelAndPromotionShareOneOutcome(t *testing.T) { - s, pool := testStore(t) - writer := executionWriter(t, s) - other, _ := testStore(t) - tenant, session := environmentInputSession(t, s) - ctx := context.Background() - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - start := make(chan struct{}) - results := make(chan sessions.EnvironmentInputReservation, 2) - errs := make(chan error, 2) - for _, finish := range []func(context.Context, string, string, string) (sessions.EnvironmentInputReservation, error){ - writer.PromoteEnvironmentInput, other.CancelEnvironmentInput, - } { - go func() { - <-start - got, err := finish(ctx, tenant, session.ID, pending.ID) - results <- got - errs <- err - }() - } - close(start) - first, second := <-results, <-results - for range 2 { - if err := <-errs; err != nil { - t.Fatal(err) - } - } - if first.State != second.State || (first.State != sessions.EnvironmentInputAdmitted && first.State != sessions.EnvironmentInputCancelled) { - t.Fatal("competing settlements diverged", first.State, second.State) - } - turns, inputs := 0, 0 - if first.State == sessions.EnvironmentInputAdmitted { - turns, inputs = 1, 2 - } - environmentInputHistory(t, pool, session.ID, turns, inputs) -} - func TestEnvironmentInputDeletionSettlesPendingAndFencesPromotion(t *testing.T) { for _, concurrent := range []bool{false, true} { t.Run(map[bool]string{false: "pending", true: "racing-promotion"}[concurrent], func(t *testing.T) { @@ -231,7 +19,7 @@ func TestEnvironmentInputDeletionSettlesPendingAndFencesPromotion(t *testing.T) done := make(chan error, 1) if concurrent { go func() { - _, err := writer.PromoteEnvironmentInput(ctx, tenant, session.ID, pending.ID) + _, err := sessionExecution(t, writer.lease).PromoteEnvironmentInput(ctx, tenant, session.ID, pending.ID) done <- err }() } @@ -250,13 +38,17 @@ func TestEnvironmentInputDeletionSettlesPendingAndFencesPromotion(t *testing.T) } } for _, action := range []func(context.Context, string, string, string) (sessions.EnvironmentInputReservation, error){ - s.GetEnvironmentInputReservation, writer.PromoteEnvironmentInput, s.CancelEnvironmentInput, + sessionAdapter(s).GetEnvironmentInputReservation, + sessionExecution(t, writer.lease).PromoteEnvironmentInput, + func(ctx context.Context, tenant, session, id string) (sessions.EnvironmentInputReservation, error) { + return cancelEnvironmentInput(ctx, s, tenant, session, id) + }, } { if _, err := action(ctx, tenant, session.ID, pending.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted reservation remained accessible", err) } } - if _, err := s.ReserveEnvironmentInput(ctx, tenant, session.ID, "late", pending.Inputs); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sessionService(t, s).ReserveEnvironmentInput(ctx, tenant, session.ID, "late", pending.Inputs); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted Session accepted reservation", err) } var state string diff --git a/services/core/internal/store/environment_inputs.go b/services/core/internal/store/environment_inputs.go deleted file mode 100644 index c2d2d4b7f..000000000 --- a/services/core/internal/store/environment_inputs.go +++ /dev/null @@ -1,287 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "fmt" - - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -// ReserveEnvironmentInput appends to active work or reserves an idle message batch. -// The Session lock decides both paths; only promotion can create a new Turn. -func (s *Store) ReserveEnvironmentInput(ctx context.Context, tenantID, sessionID, key string, inputs []sessions.Input) (sessions.EnvironmentInputReservation, error) { - if err := sessions.ValidateInputKey(key); err != nil { - return sessions.EnvironmentInputReservation{}, err - } - batch, encoded, err := validateInitialInputs(inputs) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - tenant, err := parseID(tenantID) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - var result sessions.EnvironmentInputReservation - err = s.withEnvironmentInputSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - audit := func() error { - return auditpg.RecordWriteAudit(ctx, q, tenantID, "send_events", "session", uuid.UUID(session.Bytes).String(), "") - } - previous, err := q.FindEnvironmentInputReservation(ctx, sqlc.FindEnvironmentInputReservationParams{ - SessionID: session, IdempotencyKey: key, Batch: encoded, - }) - if err == nil { - if !previous.Matches { - return sessions.ErrIdempotencyConflict - } - result, err = settleEnvironmentInput(ctx, q, tenantID, previous.EnvironmentInputReservation, sessions.EnvironmentInputExpired) - if err == nil && (result.State == sessions.EnvironmentInputPending || result.State == sessions.EnvironmentInputAdmitted) { - return audit() - } - return err - } - if !errors.Is(err, pgx.ErrNoRows) { - return err - } - receipts, err := inputBatchReceipts(ctx, q, session, key, encoded) - if err != nil { - return err - } - if len(receipts) > 0 { - // Earlier direct admission has receipts, but never had a reservation or deadline. - result = sessions.EnvironmentInputReservation{SessionID: sessionID, State: sessions.EnvironmentInputAdmitted, Receipts: receipts} - return audit() - } - if err := sessions.CheckFileWriteGate(ctx, sessionpg.BindSession(q, tenant, session)); err != nil { - return err - } - environment, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: tenant, ID: session}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrInvalidInput - } else if err != nil { - return err - } - if environment.Environment.Status == "failed" { - if kind, err := sessions.EnvironmentType(environment.Configuration); err == nil && kind == "openai_hosted" { - return sessions.ErrHostedEnvironmentFailed - } - } - if environment.Environment.Status == "failed" || environment.Environment.Status == "expired" { - return sessions.ErrEnvironmentUnavailable - } - if err := checkEnvironmentInputGate(ctx, q, session, key, encoded); err != nil { - return err - } - if active, err := q.GetActiveTurn(ctx, session); err == nil && !active.ArtifactCaptureStarted { - result = sessions.EnvironmentInputReservation{SessionID: sessionID, State: sessions.EnvironmentInputAdmitted} - for position, input := range batch { - receipt, err := admitInput(ctx, q, tenantID, session, key, int32(position), input) - if err != nil { - return err - } - result.Receipts = append(result.Receipts, receipt) - } - return audit() - } else if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return err - } - row, err := q.CreateEnvironmentInputReservation(ctx, sqlc.CreateEnvironmentInputReservationParams{ - ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, SessionID: session, IdempotencyKey: key, Batch: encoded, - }) - if err != nil { - return err - } - result, err = environmentInputFromRow(row) - if err != nil { - return err - } - return audit() - }) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - return result, nil -} - -func (s *Store) GetEnvironmentInputReservation(ctx context.Context, tenantID, sessionID, reservationID string) (sessions.EnvironmentInputReservation, error) { - id, err := parseID(reservationID) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - var result sessions.EnvironmentInputReservation - err = s.withPublicSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - row, err := q.GetEnvironmentInputReservation(ctx, sqlc.GetEnvironmentInputReservationParams{SessionID: session, ID: id}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - result, err = environmentInputOutcome(ctx, q, row) - return err - }) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - return result, nil -} - -// PromoteEnvironmentInput admits and claims work for the retained native preparation. -func (s *Store) PromoteEnvironmentInput(ctx context.Context, tenantID, sessionID, reservationID string) (sessions.EnvironmentInputReservation, error) { - if err := s.checkExecutionAuthority(); err != nil { - return sessions.EnvironmentInputReservation{}, err - } - return s.settleEnvironmentInput(ctx, tenantID, sessionID, reservationID, sessions.EnvironmentInputAdmitted) -} - -func (s *Store) CancelEnvironmentInput(ctx context.Context, tenantID, sessionID, reservationID string) (sessions.EnvironmentInputReservation, error) { - return s.settleEnvironmentInput(ctx, tenantID, sessionID, reservationID, sessions.EnvironmentInputCancelled) -} - -// FailEnvironmentInput settles a confirmed pre-admission failure. The Session -// lock and pending-state predicate preserve cancellation and newer input. -func (s *Store) FailEnvironmentInput(ctx context.Context, tenantID, sessionID, reservationID, code string) error { - if code != "model_provider_required" && code != "runtime_preparation_failed" { - return sessions.ErrInvalidInput - } - id, err := parseID(reservationID) - if err != nil { - return err - } - return s.withEnvironmentInputSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - if err := q.ExpireEnvironmentInputReservation(ctx, sqlc.ExpireEnvironmentInputReservationParams{SessionID: session, ID: id}); err != nil { - return err - } - _, err := q.FailEnvironmentInput(ctx, sqlc.FailEnvironmentInputParams{SessionID: session, ID: id, FailureCode: pgtype.Text{String: code, Valid: true}}) - return err - }) -} - -func (s *Store) ExpireEnvironmentInput(ctx context.Context, tenantID, sessionID, reservationID string) (sessions.EnvironmentInputReservation, error) { - return s.settleEnvironmentInput(ctx, tenantID, sessionID, reservationID, sessions.EnvironmentInputExpired) -} - -func (s *Store) settleEnvironmentInput(ctx context.Context, tenantID, sessionID, reservationID, state string) (sessions.EnvironmentInputReservation, error) { - id, err := parseID(reservationID) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - var result sessions.EnvironmentInputReservation - err = s.withEnvironmentInputSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - row, err := q.GetEnvironmentInputReservation(ctx, sqlc.GetEnvironmentInputReservationParams{SessionID: session, ID: id}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - result, err = settleEnvironmentInput(ctx, q, tenantID, row, state) - return err - }) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - return result, nil -} - -func settleEnvironmentInput(ctx context.Context, q *sqlc.Queries, tenantID string, row sqlc.EnvironmentInputReservation, state string) (sessions.EnvironmentInputReservation, error) { - if row.State != sessions.EnvironmentInputPending { - return environmentInputOutcome(ctx, q, row) - } - if err := q.ExpireEnvironmentInputReservation(ctx, sqlc.ExpireEnvironmentInputReservationParams{SessionID: row.SessionID, ID: row.ID}); err != nil { - return sessions.EnvironmentInputReservation{}, err - } - row, err := q.GetEnvironmentInputReservation(ctx, sqlc.GetEnvironmentInputReservationParams{SessionID: row.SessionID, ID: row.ID}) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - // Terminal outcomes are successful storage results so settlement is not rolled back. - if row.State != sessions.EnvironmentInputPending || state == sessions.EnvironmentInputExpired { - return environmentInputOutcome(ctx, q, row) - } - result, err := environmentInputFromRow(row) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - if state == sessions.EnvironmentInputAdmitted { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - if err := sessions.CheckInputStart(ctx, sessionpg.BindSession(q, tenant, row.SessionID)); err != nil { - return sessions.EnvironmentInputReservation{}, err - } - for position, input := range result.Inputs { - receipt, err := admitInput(ctx, q, tenantID, row.SessionID, row.IdempotencyKey, int32(position), input) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - result.Receipts = append(result.Receipts, receipt) - } - } - row, err = q.SettleEnvironmentInputReservation(ctx, sqlc.SettleEnvironmentInputReservationParams{SessionID: row.SessionID, ID: row.ID, State: state}) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - if state == sessions.EnvironmentInputAdmitted { - tenant, err := parseID(tenantID) - if err != nil { - return sessions.EnvironmentInputReservation{}, err - } - start := sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress} - if _, err := sessions.TransitionTurn(ctx, sessionpg.BindSession(q, tenant, row.SessionID), result.Receipts[0].TurnID, start); err != nil { - return sessions.EnvironmentInputReservation{}, err - } - } - result.State = row.State - result.SettledAt = &row.SettledAt.Time - return result, nil -} - -func environmentInputOutcome(ctx context.Context, q *sqlc.Queries, row sqlc.EnvironmentInputReservation) (sessions.EnvironmentInputReservation, error) { - result, err := environmentInputFromRow(row) - if err != nil || row.State != sessions.EnvironmentInputAdmitted { - return result, err - } - result.Receipts, err = inputBatchReceipts(ctx, q, row.SessionID, row.IdempotencyKey, row.Batch) - if err == nil && len(result.Receipts) == 0 { - err = errors.New("admitted Environment input has no receipts") - } - return result, err -} - -func environmentInputFromRow(row sqlc.EnvironmentInputReservation) (sessions.EnvironmentInputReservation, error) { - result := sessions.EnvironmentInputReservation{ - ID: uuid.UUID(row.ID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), - State: row.State, IsInitial: row.IsInitial, CreatedAt: row.CreatedAt.Time, Deadline: row.Deadline.Time, - } - if row.SettledAt.Valid { - result.SettledAt = &row.SettledAt.Time - } - if err := json.Unmarshal(row.Batch, &result.Inputs); err != nil { - return sessions.EnvironmentInputReservation{}, fmt.Errorf("decode Environment input: %w", err) - } - return result, nil -} - -func checkEnvironmentInputGate(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, key string, batch json.RawMessage) error { - gate, err := q.CheckEnvironmentInputGate(ctx, sqlc.CheckEnvironmentInputGateParams{SessionID: session, IdempotencyKey: key, Batch: batch}) - if err != nil { - return err - } - if !gate.Matches { - return sessions.ErrIdempotencyConflict - } - if gate.Blocked { - return sessions.ErrInputPending - } - return nil -} diff --git a/services/core/internal/store/environment_inputs_test.go b/services/core/internal/store/environment_inputs_test.go deleted file mode 100644 index 2b2074e70..000000000 --- a/services/core/internal/store/environment_inputs_test.go +++ /dev/null @@ -1,274 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "reflect" - "sync" - "testing" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" - "github.com/jackc/pgx/v5/pgxpool" -) - -func environmentInputSession(t *testing.T, s *Store) (string, sessions.Session) { - t.Helper() - tenant := uuid.NewString() - session, err := s.CreateSession(context.Background(), tenant, environmentInput("session", "self_hosted", "/workspace")) - if err != nil { - t.Fatal(err) - } - return tenant, session -} - -func reserveEnvironmentInput(t *testing.T, s *Store, tenant, session, key string) sessions.EnvironmentInputReservation { - t.Helper() - got, err := s.ReserveEnvironmentInput(context.Background(), tenant, session, key, []sessions.Input{messageInput("first"), messageInput("second")}) - if err != nil { - t.Fatal(err) - } - return got -} - -func environmentInputHistory(t *testing.T, pool *pgxpool.Pool, session string, turns, inputs int) { - t.Helper() - var gotTurns, gotInputs, items, events int - err := pool.QueryRow(context.Background(), ` - SELECT (SELECT count(*) FROM turns WHERE session_id=$1), - (SELECT count(*) FROM turn_inputs WHERE session_id=$1), - (SELECT count(*) FROM session_items WHERE session_id=$1), - (SELECT count(*) FROM session_events WHERE session_id=$1 - AND (payload ? 'turn' OR payload->'event' ? 'item'))`, session).Scan(&gotTurns, &gotInputs, &items, &events) - if err != nil || gotTurns != turns || gotInputs != inputs || items != inputs || (inputs == 0 && events != 0) { - t.Fatal("history", gotTurns, gotInputs, items, events, err) - } -} - -func TestEnvironmentInputReservationConcurrentIdentity(t *testing.T) { - s, pool := testStore(t) - other, _ := testStore(t) - tenant, session := environmentInputSession(t, s) - ctx := context.Background() - batch := []sessions.Input{ - {Kind: "message", Payload: json.RawMessage(`{"text":"first","detail":{"a":1,"b":2}}`)}, - messageInput("second"), - } - const count = 8 - results := make(chan sessions.EnvironmentInputReservation, count) - var wg sync.WaitGroup - for i := range count { - wg.Add(1) - go func() { - defer wg.Done() - st := s - inputs := append([]sessions.Input(nil), batch...) - if i%2 == 0 { - st = other - inputs[0].Payload = json.RawMessage(` { "detail": {"b": 2, "a": 1}, "text": "first" } `) - } - got, err := st.ReserveEnvironmentInput(ctx, tenant, session.ID, "request", inputs) - if err != nil { - t.Error(err) - return - } - results <- got - }() - } - wg.Wait() - close(results) - var first sessions.EnvironmentInputReservation - received := 0 - for result := range results { - received++ - if first.ID == "" { - first = result - } - if !reflect.DeepEqual(first, result) { - t.Fatal("reservation identity changed", first, result) - } - } - if received != count || first.State != sessions.EnvironmentInputPending || first.ID == "" || first.Deadline.Sub(first.CreatedAt) != 5*time.Minute || first.SettledAt != nil || len(first.Receipts) != 0 { - t.Fatal("invalid pending result", received, first) - } - environmentInputHistory(t, pool, session.ID, 0, 0) - for _, changed := range [][]sessions.Input{batch[:1], {batch[1], batch[0]}, {messageInput("changed"), batch[1]}} { - if _, err := s.ReserveEnvironmentInput(ctx, tenant, session.ID, "request", changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { - t.Fatal("changed request accepted", err) - } - } - if _, err := other.ReserveEnvironmentInput(ctx, tenant, session.ID, "other", batch); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal("second pending request accepted", err) - } - pool.Close() - restarted, _ := testStore(t) - got, err := restarted.ReserveEnvironmentInput(ctx, tenant, session.ID, "request", batch) - if err != nil || !reflect.DeepEqual(first, got) { - t.Fatal("restart changed deadline or identity", got, err) - } -} - -func TestEnvironmentInputReservationPromotionAndDirectRetries(t *testing.T) { - s, pool := testStore(t) - writer := executionWriter(t, s) - tenant, session := environmentInputSession(t, s) - ctx := context.Background() - first := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - for _, request := range []struct { - key string - inputs []sessions.Input - want error - }{ - {"pending", first.Inputs, sessions.ErrTurnConflict}, - {"pending", []sessions.Input{messageInput("changed")}, sessions.ErrIdempotencyConflict}, - {"later", []sessions.Input{messageInput("later")}, sessions.ErrTurnConflict}, - {"cancel", []sessions.Input{{Kind: "cancel", Payload: json.RawMessage(`{}`)}}, sessions.ErrTurnConflict}, - } { - if _, err := s.SubmitInputs(ctx, tenant, session.ID, request.key, request.inputs); !errors.Is(err, request.want) { - t.Fatal("direct path bypassed reservation", request.key, err) - } - } - environmentInputHistory(t, pool, session.ID, 0, 0) - promoted, err := writer.PromoteEnvironmentInput(ctx, tenant, session.ID, first.ID) - if err != nil || promoted.State != sessions.EnvironmentInputAdmitted || promoted.SettledAt == nil || len(promoted.Receipts) != 2 || !promoted.Deadline.Equal(first.Deadline) { - t.Fatal(promoted, err) - } - for i, receipt := range promoted.Receipts { - if receipt.Replayed || receipt.TurnID == "" || receipt.TurnID != promoted.Receipts[0].TurnID || (i > 0 && receipt.Sequence <= promoted.Receipts[i-1].Sequence) { - t.Fatal("promotion receipts", promoted.Receipts) - } - } - environmentInputHistory(t, pool, session.ID, 1, 2) - for _, read := range []func() (sessions.EnvironmentInputReservation, error){ - func() (sessions.EnvironmentInputReservation, error) { - return writer.PromoteEnvironmentInput(ctx, tenant, session.ID, first.ID) - }, - func() (sessions.EnvironmentInputReservation, error) { - return s.GetEnvironmentInputReservation(ctx, tenant, session.ID, first.ID) - }, - func() (sessions.EnvironmentInputReservation, error) { - return s.ReserveEnvironmentInput(ctx, tenant, session.ID, "pending", first.Inputs) - }, - } { - retry, err := read() - if err != nil || retry.ID != first.ID || !retry.Deadline.Equal(first.Deadline) || retry.State != sessions.EnvironmentInputAdmitted || len(retry.Receipts) != 2 { - t.Fatal(retry, err) - } - for i, receipt := range retry.Receipts { - if !receipt.Replayed || receipt.Sequence != promoted.Receipts[i].Sequence || receipt.TurnID != promoted.Receipts[i].TurnID { - t.Fatal("retry changed admission", receipt) - } - } - } - retry, err := s.SubmitInputs(ctx, tenant, session.ID, "pending", first.Inputs) - if err != nil || len(retry) != 2 || !retry[0].Replayed || retry[0].Sequence != promoted.Receipts[0].Sequence { - t.Fatal("direct retry after promotion", retry, err) - } - awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, writer.pool) - if err := writer.lease.Close(ctx); err != nil { - t.Fatal(err) - } - awaitRelease() - pool.Close() - restarted, pool := testStore(t) - after, err := executionWriter(t, restarted).PromoteEnvironmentInput(ctx, tenant, session.ID, first.ID) - if err != nil || after.State != sessions.EnvironmentInputAdmitted || after.Receipts[0].Sequence != promoted.Receipts[0].Sequence { - t.Fatal("restart repeated promotion", after, err) - } - environmentInputHistory(t, pool, session.ID, 1, 2) -} - -func TestEnvironmentInputReservationKeepsEarlierDirectIdentity(t *testing.T) { - s, _ := testStore(t) - tenant, session := environmentInputSession(t, s) - ctx := context.Background() - input := messageInput("already admitted") - receipts, err := s.SubmitInputs(ctx, tenant, session.ID, "direct", []sessions.Input{input}) - if err != nil { - t.Fatal(err) - } - transition(t, s, tenant, session.ID, receipts[0].TurnID, sessions.TurnQueued, sessions.TurnInProgress) - transition(t, s, tenant, session.ID, receipts[0].TurnID, sessions.TurnInProgress, sessions.TurnCompleted) - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "new") - got, err := s.ReserveEnvironmentInput(ctx, tenant, session.ID, "direct", []sessions.Input{input}) - if err != nil || got.State != sessions.EnvironmentInputAdmitted || got.ID != "" || !got.Deadline.IsZero() || len(got.Receipts) != 1 || got.Receipts[0].Sequence != receipts[0].Sequence { - t.Fatal("direct admission gained a reservation", got, err) - } - if _, err := s.ReserveEnvironmentInput(ctx, tenant, session.ID, "direct", []sessions.Input{messageInput("changed")}); !errors.Is(err, sessions.ErrIdempotencyConflict) { - t.Fatal(err) - } - retry, err := s.SubmitInputs(ctx, tenant, session.ID, "direct", []sessions.Input{input}) - if err != nil || len(retry) != 1 || !retry[0].Replayed { - t.Fatal(retry, err) - } - retained, err := s.GetEnvironmentInputReservation(ctx, tenant, session.ID, pending.ID) - if err != nil || !reflect.DeepEqual(retained, pending) { - t.Fatal("old retry affected new pending input", retained, err) - } -} - -func TestEnvironmentInputReservationRejectsUnsupportedOrForeignState(t *testing.T) { - s, _ := testStore(t) - writer := executionWriter(t, s) - tenant, session := environmentInputSession(t, s) - ctx := context.Background() - pending := reserveEnvironmentInput(t, s, tenant, session.ID, "pending") - for _, read := range []func() error{ - func() error { - _, err := s.GetEnvironmentInputReservation(ctx, uuid.NewString(), session.ID, pending.ID) - return err - }, - func() error { - _, err := writer.PromoteEnvironmentInput(ctx, uuid.NewString(), session.ID, pending.ID) - return err - }, - func() error { - _, err := s.CancelEnvironmentInput(ctx, uuid.NewString(), session.ID, pending.ID) - return err - }, - func() error { - _, err := s.ReserveEnvironmentInput(ctx, uuid.NewString(), session.ID, "new", pending.Inputs) - return err - }, - } { - if err := read(); !errors.Is(err, sessions.ErrNotFound) { - t.Fatal("foreign access", err) - } - } - other, err := s.CreateSession(ctx, tenant, environmentInput("other", "self_hosted", "/workspace")) - if err != nil { - t.Fatal(err) - } - for _, action := range []func(context.Context, string, string, string) (sessions.EnvironmentInputReservation, error){ - s.GetEnvironmentInputReservation, writer.PromoteEnvironmentInput, s.CancelEnvironmentInput, s.ExpireEnvironmentInput, - } { - if _, err := action(ctx, tenant, other.ID, pending.ID); !errors.Is(err, sessions.ErrNotFound) { - t.Fatal("reservation crossed Session ownership", err) - } - if _, err := action(ctx, uuid.NewString(), session.ID, pending.ID); !errors.Is(err, sessions.ErrNotFound) { - t.Fatal("reservation crossed tenant ownership", err) - } - } - retained, err := s.GetEnvironmentInputReservation(ctx, tenant, session.ID, pending.ID) - if err != nil || !reflect.DeepEqual(retained, pending) { - t.Fatal("foreign operations changed reservation", retained, err) - } - for _, invalid := range [][]sessions.Input{nil, {{Kind: "cancel", Payload: json.RawMessage(`{}`)}}, {{Kind: "tool_result", Payload: json.RawMessage(`{}`)}}} { - if _, err := s.ReserveEnvironmentInput(ctx, tenant, session.ID, "invalid", invalid); !errors.Is(err, sessions.ErrInvalidInput) { - t.Fatal("unsupported reservation", err) - } - } - noneTenant, none := newTurnSession(t, s) - if _, err := s.ReserveEnvironmentInput(ctx, noneTenant, none.ID, "none", pending.Inputs); !errors.Is(err, sessions.ErrInvalidInput) { - t.Fatal("none reservation", err) - } - activeTenant, active := environmentInputSession(t, s) - activeInput := submitMessage(t, s, activeTenant, active.ID, "active") - steer, err := s.ReserveEnvironmentInput(ctx, activeTenant, active.ID, "new", pending.Inputs) - if err != nil || steer.State != sessions.EnvironmentInputAdmitted || steer.ID != "" || !steer.Deadline.IsZero() || len(steer.Receipts) != len(pending.Inputs) || steer.Receipts[0].TurnID != activeInput.TurnID { - t.Fatal("active input did not retain the existing Turn", steer, err) - } -} diff --git a/services/core/internal/store/environment_steering_order_test.go b/services/core/internal/store/environment_steering_order_test.go deleted file mode 100644 index 94cfcf094..000000000 --- a/services/core/internal/store/environment_steering_order_test.go +++ /dev/null @@ -1,104 +0,0 @@ -package store - -import ( - "context" - "errors" - "testing" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -func TestEnvironmentActiveInputSerializesWithCompletion(t *testing.T) { - for _, completionFirst := range []bool{false, true} { - name := "input-first" - if completionFirst { - name = "completion-first" - } - t.Run(name, func(t *testing.T) { - s, pool := testStore(t) - writer := executionWriter(t, s) - tenant, session := environmentInputSession(t, s) - original := submitMessage(t, s, tenant, session.ID, "original") - transition(t, s, tenant, session.ID, original.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) - defer cancel() - blocker, err := pool.Begin(ctx) - if err != nil { - t.Fatal(err) - } - defer func() { _ = blocker.Rollback(context.Background()) }() - var blockerPID int32 - if err := blocker.QueryRow(ctx, "SELECT pg_backend_pid() FROM sessions WHERE id=$1 FOR UPDATE", session.ID).Scan(&blockerPID); err != nil { - t.Fatal(err) - } - waiter := func(pid int32) int32 { - t.Helper() - for ctx.Err() == nil { - var blocked int32 - if err := pool.QueryRow(ctx, "SELECT pid FROM pg_stat_activity WHERE $1=ANY(pg_blocking_pids(pid)) ORDER BY pid LIMIT 1", pid).Scan(&blocked); err == nil { - return blocked - } - time.Sleep(5 * time.Millisecond) - } - t.Fatal("transaction did not wait for the Session lock") - return 0 - } - type admission struct { - value sessions.EnvironmentInputReservation - err error - } - admitted := make(chan admission, 1) - completed := make(chan error, 1) - batch := []sessions.Input{messageInput("first"), messageInput("second")} - input := func() { - value, err := s.ReserveEnvironmentInput(ctx, tenant, session.ID, "racing-input", batch) - admitted <- admission{value, err} - } - complete := func() { - _, err := completeExecution(ctx, t, writer, tenant, session.ID, original.TurnID, sessions.TurnCompleted, nil, "", original.Sequence) - completed <- err - } - first, second := input, complete - if completionFirst { - first, second = complete, input - } - go first() - firstPID := waiter(blockerPID) - go second() - waiter(firstPID) - if err := blocker.Commit(ctx); err != nil { - t.Fatal(err) - } - got, completionErr := <-admitted, <-completed - if got.err != nil { - t.Fatal(got.err) - } - if completionFirst { - if completionErr != nil || got.value.State != sessions.EnvironmentInputPending || got.value.ID == "" || len(got.value.Receipts) != 0 || got.value.Deadline.Sub(got.value.CreatedAt) != 5*time.Minute { - t.Fatal("completion winner did not leave new input waiting for preparation", completionErr, got.value) - } - environmentInputHistory(t, pool, session.ID, 1, 1) - prepared, err := writer.PromoteEnvironmentInput(ctx, tenant, session.ID, got.value.ID) - if err != nil || len(prepared.Receipts) != 2 || prepared.Receipts[0].TurnID == original.TurnID { - t.Fatal("prepared successor reused terminal work", err) - } - retry, err := s.ReserveEnvironmentInput(ctx, tenant, session.ID, "racing-input", batch) - if err != nil || retry.ID != got.value.ID || !retry.Deadline.Equal(got.value.Deadline) || len(retry.Receipts) != 2 || !retry.Receipts[0].Replayed || retry.Receipts[0].TurnID != prepared.Receipts[0].TurnID { - t.Fatal("active retry replaced its original reservation", retry, err) - } - if _, err := completeExecution(ctx, t, writer, tenant, session.ID, prepared.Receipts[0].TurnID, sessions.TurnCompleted, nil, "", prepared.Receipts[1].Sequence); err != nil { - t.Fatal(err) - } - } else { - if !errors.Is(completionErr, sessions.ErrUnappliedInputs) || got.value.State != sessions.EnvironmentInputAdmitted || got.value.ID != "" || !got.value.Deadline.IsZero() || len(got.value.Receipts) != 2 || got.value.Receipts[0].TurnID != original.TurnID || got.value.Receipts[0].Replayed { - t.Fatal("admitted input escaped the original Turn or application fence", completionErr, got.value) - } - environmentInputHistory(t, pool, session.ID, 1, 3) - if _, err := completeExecution(ctx, t, writer, tenant, session.ID, original.TurnID, sessions.TurnCompleted, nil, "", got.value.Receipts[1].Sequence); err != nil { - t.Fatal("completion after controlled application failed", err) - } - } - }) - } -} diff --git a/services/core/internal/store/environment_work_test.go b/services/core/internal/store/environment_work_test.go index 5085d3585..0c07a6321 100644 --- a/services/core/internal/store/environment_work_test.go +++ b/services/core/internal/store/environment_work_test.go @@ -23,7 +23,7 @@ func TestEnvironmentInputWorkFiltersAndPagesDevices(t *testing.T) { case "expired": makeEnvironmentExpiryDue(t, pool, &pending) case "cancelled": - if _, err := h.s.CancelEnvironmentInput(t.Context(), h.tenant, pending.SessionID, pending.ID); err != nil { + if _, err := store.CancelEnvironmentInput(t.Context(), h.s, h.tenant, pending.SessionID, pending.ID); err != nil { t.Fatal(err) } case "deleted": @@ -44,7 +44,7 @@ func TestEnvironmentInputWorkFiltersAndPagesDevices(t *testing.T) { if err := fixtureSessionService(t, h.db).RevokeDevice(t.Context(), h.tenant, runtime.device.ID); err != nil { t.Fatal(err) } - work, err := h.s.ListEnvironmentInputWork(t.Context(), "", []string{runtime.device.ID}) + work, err := store.SessionAdapter(h.s).ListEnvironmentInputWork(t.Context(), "", []string{runtime.device.ID}) if err != nil || len(work) != 0 { t.Fatal("revoked Runtime selected", work, err) } @@ -55,7 +55,7 @@ func TestEnvironmentInputWorkFiltersAndPagesDevices(t *testing.T) { } } for _, devices := range [][]string{nil, {}, {uuid.NewString()}} { - work, err := h.s.ListEnvironmentInputWork(t.Context(), "", devices) + work, err := store.SessionAdapter(h.s).ListEnvironmentInputWork(t.Context(), "", devices) if err != nil || len(work) != 0 { t.Fatal("unconnected work selected", work, err) } @@ -67,7 +67,7 @@ func TestEnvironmentInputWorkFiltersAndPagesDevices(t *testing.T) { for _, runtime := range h.environments { devices = append(devices, runtime.device.ID) } - work, err := h.s.ListEnvironmentInputWork(t.Context(), cursor, devices) + work, err := store.SessionAdapter(h.s).ListEnvironmentInputWork(t.Context(), cursor, devices) if err != nil || len(work) != count { t.Fatal("environment work page", len(work), count, err) } diff --git a/services/core/internal/store/environment_worker_helpers_test.go b/services/core/internal/store/environment_worker_helpers_test.go index d39561867..4ea8d7485 100644 --- a/services/core/internal/store/environment_worker_helpers_test.go +++ b/services/core/internal/store/environment_worker_helpers_test.go @@ -48,7 +48,7 @@ func unboundWorkerEnvironmentReservation(t *testing.T, h *dispatchHarness) sessi if err != nil { t.Fatal(err) } - pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, session.ID, "work", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) + pending, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, session.ID, "work", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) if err != nil { t.Fatal(err) } @@ -115,7 +115,7 @@ func awaitWorkerEnvironmentRun(t *testing.T, ctx context.Context, s *store.Store var run execution.EnvironmentRun awaitDaemonRemoteCondition(t, ctx, 5*time.Minute, "worker terminal Environment Turn", func() bool { var err error - run.Reservation, err = s.GetEnvironmentInputReservation(ctx, tenant, pending.SessionID, pending.ID) + run.Reservation, err = store.SessionAdapter(s).GetEnvironmentInputReservation(ctx, tenant, pending.SessionID, pending.ID) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/environment_worker_scan_test.go b/services/core/internal/store/environment_worker_scan_test.go index f160dbb04..56860b921 100644 --- a/services/core/internal/store/environment_worker_scan_test.go +++ b/services/core/internal/store/environment_worker_scan_test.go @@ -44,7 +44,7 @@ func TestWorkerEnvironmentRetriesNewlyReadyAtNextScan(t *testing.T) { runtime.write(prepare.ID, proto.TypePreparationStatus, proto.PreparationStatusPayload{Handle: handle, Revision: 2, State: "failed"}) nextWorkerFrame(t, frames, proto.TypeExecutionRelease) stop() - stored, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, pending.SessionID, pending.ID) + stored, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, pending.SessionID, pending.ID) if err != nil || stored.State != sessions.EnvironmentInputPending || !stored.Deadline.Equal(pending.Deadline) || len(stored.Receipts) != 0 { t.Fatal("readiness retry changed pending identity or admitted work", stored, err) } diff --git a/services/core/internal/store/environment_worker_test.go b/services/core/internal/store/environment_worker_test.go index 275e1220f..58d851208 100644 --- a/services/core/internal/store/environment_worker_test.go +++ b/services/core/internal/store/environment_worker_test.go @@ -82,7 +82,7 @@ func TestWorkerEnvironmentSharesCapacityThroughClaimAndCleanup(t *testing.T) { tenant, due := newEnvironmentExpiryReservation(t, h.s) makeEnvironmentExpiryDue(t, pool, &due) waitEnvironmentExpiry(t, h.s, tenant, due) - if _, err := h.s.CancelEnvironmentInput(t.Context(), h.tenant, waiting.SessionID, waiting.ID); err != nil { + if _, err := store.CancelEnvironmentInput(t.Context(), h.s, h.tenant, waiting.SessionID, waiting.ID); err != nil { t.Fatal(err) } // Distinct Runtime sockets do not promise cross-socket delivery order. @@ -159,7 +159,7 @@ func TestWorkerEnvironmentRetriesPendingWithoutExtendingDeadline(t *testing.T) { } stop() nextWorkerFrame(t, frames, proto.TypeExecutionRelease) - stored, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, pending.SessionID, pending.ID) + stored, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, pending.SessionID, pending.ID) if err != nil || stored.State != sessions.EnvironmentInputPending || !stored.Deadline.Equal(pending.Deadline) || len(stored.Receipts) != 0 { t.Fatal("retry or shutdown changed the original reservation", stored, err) } diff --git a/services/core/internal/store/execution.go b/services/core/internal/store/execution.go index 0978511c5..92d5f5fc9 100644 --- a/services/core/internal/store/execution.go +++ b/services/core/internal/store/execution.go @@ -1,13 +1,6 @@ package store -import ( - "errors" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" -) - -// ErrExecutionAuthority rejects an execution-only operation on a pooled Store. -var ErrExecutionAuthority = errors.New("operation requires the execution writer") +import "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" // NewExecution returns the execution writer built on lease, which the caller // acquired and closes. The writer's Session and execution-only transactions run @@ -18,12 +11,3 @@ func NewExecution(s *Store, lease *pgunit.Lease) *Store { writer.writer, writer.lease = lease, lease return &writer } - -// checkExecutionAuthority only validates. The connection was fixed when the -// Store was constructed. -func (s *Store) checkExecutionAuthority() error { - if s.lease == nil { - return ErrExecutionAuthority - } - return nil -} diff --git a/services/core/internal/store/execution_events_test.go b/services/core/internal/store/execution_events_test.go index 2ba74d5c9..5e8b963e9 100644 --- a/services/core/internal/store/execution_events_test.go +++ b/services/core/internal/store/execution_events_test.go @@ -39,7 +39,7 @@ func TestExecutionPersistsLiveAndCancelledPartialOutput(t *testing.T) { } time.Sleep(20 * time.Millisecond) } - if _, err := h.s.RequestCancel(ctx, h.tenant, h.session.ID, "stop"); err != nil { + if _, err := store.RequestCancel(ctx, h.s, h.tenant, h.session.ID, "stop"); err != nil { t.Fatal(err) } cancelEnv := h.read(proto.TypePromptCancel) diff --git a/services/core/internal/store/execution_messages_test.go b/services/core/internal/store/execution_messages_test.go index ccb8f31eb..3df7562ee 100644 --- a/services/core/internal/store/execution_messages_test.go +++ b/services/core/internal/store/execution_messages_test.go @@ -54,7 +54,7 @@ func TestExecutionNegotiatesAndPersistsMessageObservations(t *testing.T) { h.write(input.TurnID, proto.TypeOutputMessage, proto.OutputMessagePayload{ID: "a", Status: "completed", Phase: "commentary", Text: &text}) h.write(input.TurnID, proto.TypeOutputMessage, proto.OutputMessagePayload{ID: "b", Status: "in_progress", Phase: "final_answer"}) h.write(input.TurnID, proto.TypeDelta, proto.DeltaPayload{ItemID: "b", Delta: "partial", Sequence: 2}) - if _, err := h.s.RequestCancel(ctx, h.tenant, h.session.ID, "cancel"); err != nil { + if _, err := store.RequestCancel(ctx, h.s, h.tenant, h.session.ID, "cancel"); err != nil { t.Fatal(err) } env = h.read(proto.TypePromptCancel) diff --git a/services/core/internal/store/execution_test.go b/services/core/internal/store/execution_test.go index c50428880..8104027e5 100644 --- a/services/core/internal/store/execution_test.go +++ b/services/core/internal/store/execution_test.go @@ -169,7 +169,7 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { mustReject("completion", err) _, err = operations.TransitionTurn(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed}) mustReject("reconciliation", err) - _, err = writer.ExpireEnvironmentInputs(t.Context()) + _, err = operations.ExpireEnvironmentInputs(t.Context()) mustReject("input expiry", err) mustReject("ownership check", writer.lease.CheckOwnership(t.Context())) after, err := sessionAdapter(s).GetSession(t.Context(), tenant, active.ID) @@ -276,7 +276,7 @@ func TestExecutionWriterSerializesWritesOnItsLease(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), time.Second) defer cancel() task := tasks[0] - _, err := s.SubmitMessage(ctx, task.tenant, task.session, "public", json.RawMessage(`{"text":"additional"}`)) + _, err := sendMessage(ctx, s, task.tenant, task.session, "public", json.RawMessage(`{"text":"additional"}`)) close(release) if err != nil { t.Fatal("public admission used owner gate", err) @@ -285,27 +285,3 @@ func TestExecutionWriterSerializesWritesOnItsLease(t *testing.T) { t.Fatal(err) } } - -func TestPooledStoreHasNoExecutionAuthority(t *testing.T) { - s, _ := testStore(t) - tenant, session := newTurnSession(t, s) - submitMessage(t, s, tenant, session.ID, "start") - before, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) - if err != nil { - t.Fatal(err) - } - cursor, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) - if err != nil { - t.Fatal(err) - } - if _, err := s.ExpireEnvironmentInputs(t.Context()); !errors.Is(err, ErrExecutionAuthority) { - t.Fatalf("pooled Store ran input expiry: %v", err) - } - after, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) - if err != nil || !reflect.DeepEqual(before, after) { - t.Fatal("rejected execution operation changed the Session", after, err) - } - if next, err := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID); err != nil || next != cursor { - t.Fatal("rejected execution operation published events", next, err) - } -} diff --git a/services/core/internal/store/execution_tools_test.go b/services/core/internal/store/execution_tools_test.go index fe1484b4d..2db63f908 100644 --- a/services/core/internal/store/execution_tools_test.go +++ b/services/core/internal/store/execution_tools_test.go @@ -40,7 +40,7 @@ func TestExecutionNegotiatesAndPersistsToolObservations(t *testing.T) { } h.write(input.TurnID, proto.TypeToolCall, proto.ToolCallPayload{ID: id, Stage: stage, Observation: &observation}) } - if _, err := h.s.RequestCancel(ctx, h.tenant, h.session.ID, "cancel"); err != nil { + if _, err := store.RequestCancel(ctx, h.s, h.tenant, h.session.ID, "cancel"); err != nil { t.Fatal(err) } env = h.read(proto.TypePromptCancel) diff --git a/services/core/internal/store/export_test.go b/services/core/internal/store/export_test.go index 6e3040cc3..50d908034 100644 --- a/services/core/internal/store/export_test.go +++ b/services/core/internal/store/export_test.go @@ -28,6 +28,27 @@ func SessionAdapter(s *Store) *sessionpg.Store { return sessionAdapter(s) } // SessionService is the Session service on SessionAdapter(s). func SessionService(t testing.TB, s *Store) *sessions.Service { return sessionService(t, s) } +// SubmitInputs admits inputs through the Session service on s's database. +func SubmitInputs(ctx context.Context, s *Store, tenant, session, key string, inputs []sessions.Input) ([]sessions.InputReceipt, error) { + return submitInputs(ctx, s, tenant, session, key, inputs) +} + +// SendMessage admits one message input with payload. +func SendMessage(ctx context.Context, s *Store, tenant, session, key string, payload json.RawMessage) (sessions.InputReceipt, error) { + return sendMessage(ctx, s, tenant, session, key, payload) +} + +// RequestCancel admits one cancel input. +func RequestCancel(ctx context.Context, s *Store, tenant, session, key string) (sessions.InputReceipt, error) { + return requestCancel(ctx, s, tenant, session, key) +} + +// CancelEnvironmentInput cancels the Session's pending Environment input as +// Session cancellation does, then reads the reservation back. +func CancelEnvironmentInput(ctx context.Context, s *Store, tenant, session, reservation string) (sessions.EnvironmentInputReservation, error) { + return cancelEnvironmentInput(ctx, s, tenant, session, reservation) +} + // TransitionTurn moves the Turn as the execution owner does, in a Session // transaction on s's writer. func TransitionTurn(ctx context.Context, s *Store, tenant, session, turn string, transition sessions.TurnTransition) (sessions.Turn, error) { @@ -144,6 +165,6 @@ func SubmitFixtureFunctionResult(ctx context.Context, s *Store, tenantID, sessio if err != nil { return err } - _, err = s.SubmitInputs(ctx, tenantID, sessionID, uuid.NewString(), []sessions.Input{{Kind: "tool_result", Payload: payload}}) + _, err = SubmitInputs(ctx, s, tenantID, sessionID, uuid.NewString(), []sessions.Input{{Kind: "tool_result", Payload: payload}}) return err } diff --git a/services/core/internal/store/function_execution_native_test.go b/services/core/internal/store/function_execution_native_test.go index a6ed06c0a..665b98188 100644 --- a/services/core/internal/store/function_execution_native_test.go +++ b/services/core/internal/store/function_execution_native_test.go @@ -32,7 +32,7 @@ func TestNativeFunctionExecutionPersistsCallsResultsAndContinuity(t *testing.T) t.Fatal(action, state.LastTurn) } if index == 2 { - if _, err := h.s.RequestCancel(ctx, h.tenant, h.session.ID, "native-cancel"); err != nil { + if _, err := store.RequestCancel(ctx, h.s, h.tenant, h.session.ID, "native-cancel"); err != nil { t.Fatal(err) } h.finished(running, sessions.TurnCancelled) diff --git a/services/core/internal/store/function_execution_test.go b/services/core/internal/store/function_execution_test.go index 0aca0f12c..95e6c0835 100644 --- a/services/core/internal/store/function_execution_test.go +++ b/services/core/internal/store/function_execution_test.go @@ -137,7 +137,7 @@ func TestExecutionFunctionsCancellationAndUnconfirmedResults(t *testing.T) { } status := sessions.TurnFailed if cancel { - if _, err := h.s.RequestCancel(t.Context(), h.tenant, h.session.ID, "cancel"); err != nil { + if _, err := store.RequestCancel(t.Context(), h.s, h.tenant, h.session.ID, "cancel"); err != nil { t.Fatal(err) } var request proto.PromptCancelPayload diff --git a/services/core/internal/store/function_images_native_test.go b/services/core/internal/store/function_images_native_test.go index 3a00438fa..54c2284eb 100644 --- a/services/core/internal/store/function_images_native_test.go +++ b/services/core/internal/store/function_images_native_test.go @@ -77,7 +77,7 @@ func TestNativeFunctionImagePublicExecution(t *testing.T) { if err != nil || !call.Applied { t.Fatal("function delivery acknowledgement missing", err) } - inputs, err := h.s.ListTurnInputs(ctx, h.tenant, proof.Session, item.Turn, 0, 100) + inputs, err := store.SessionAdapter(h.s).ListTurnInputs(ctx, h.tenant, proof.Session, item.Turn, 0, 100) if err != nil || len(inputs) != 2 || inputs[0].Kind != "message" || inputs[1].Kind != "tool_result" { t.Fatal("function result admission duplicated or mutated", err) } diff --git a/services/core/internal/store/function_input_execution_test.go b/services/core/internal/store/function_input_execution_test.go index fd0964274..5d3141eb0 100644 --- a/services/core/internal/store/function_input_execution_test.go +++ b/services/core/internal/store/function_input_execution_test.go @@ -20,11 +20,11 @@ func TestExecutionFunctionInputBatchStillSteersMessages(t *testing.T) { state := functionState(t, h, 1) raw, _ := json.Marshal(sessions.FunctionResultInput{TurnID: input.TurnID, CallID: state.RequiredActions[0].CallID, Result: json.RawMessage(`{"success":true,"output":"answer"}`)}) batch := []sessions.Input{{Kind: "tool_result", Payload: raw}, {Kind: "message", Payload: json.RawMessage(`{"text":"Follow up"}`)}} - receipts, err := h.s.SubmitInputs(t.Context(), h.tenant, h.session.ID, "mixed", batch) + receipts, err := store.SubmitInputs(t.Context(), h.s, h.tenant, h.session.ID, "mixed", batch) if err != nil { t.Fatal(err) } - if _, err := h.s.SubmitInputs(t.Context(), h.tenant, h.session.ID, "mixed", batch); err != nil { + if _, err := store.SubmitInputs(t.Context(), h.s, h.tenant, h.session.ID, "mixed", batch); err != nil { t.Fatal(err) } resultSeen, messageSeen := false, false diff --git a/services/core/internal/store/function_inputs_public_test.go b/services/core/internal/store/function_inputs_public_test.go index 7570cbe0b..4ad8bade0 100644 --- a/services/core/internal/store/function_inputs_public_test.go +++ b/services/core/internal/store/function_inputs_public_test.go @@ -34,7 +34,7 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(ctx, tenant, session.ID, "start", json.RawMessage(`{"text":"fixture"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "start", json.RawMessage(`{"text":"fixture"}`)) if err != nil { t.Fatal(err) } @@ -104,14 +104,14 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { } } } - history, err := s.ListTurnInputs(ctx, tenant, session.ID, input.TurnID, 0, 100) + history, err := store.SessionAdapter(s).ListTurnInputs(ctx, tenant, session.ID, input.TurnID, 0, 100) if err != nil || len(history) != 6 { t.Fatal(history, err) } if _, err := store.TransitionTurn(ctx, s, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnWaiting, Status: sessions.TurnFailed}); err != nil { t.Fatal(err) } - next, err := s.SubmitMessage(ctx, tenant, session.ID, "next", json.RawMessage(`{"text":"next"}`)) + next, err := store.SendMessage(ctx, s, tenant, session.ID, "next", json.RawMessage(`{"text":"next"}`)) if err != nil { t.Fatal(err) } @@ -121,7 +121,7 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { if err := command.Wait(); err != nil { t.Fatalf("SDK retry: %v %s", err, stderr.String()) } - history, err = s.ListTurnInputs(ctx, tenant, session.ID, next.TurnID, 0, 100) + history, err = store.SessionAdapter(s).ListTurnInputs(ctx, tenant, session.ID, next.TurnID, 0, 100) if err != nil || len(history) != 1 { t.Fatal(history, err) } diff --git a/services/core/internal/store/function_inputs_test.go b/services/core/internal/store/function_inputs_test.go index 5feae8e64..3a82d5ac5 100644 --- a/services/core/internal/store/function_inputs_test.go +++ b/services/core/internal/store/function_inputs_test.go @@ -40,7 +40,7 @@ func TestFunctionInputBatchesPersistAndReplayWithoutRetargeting(t *testing.T) { tenant, session, turn := functionInputFixture(t, s, functionExecution(t)) full := `{"success":false,"output":[{"type":"input_text","text":""},{"type":"input_image","image_url":"data:image/png;base64,AA=="},{"type":"input_text","text":"after"}],"error":"failed"}` batch := []sessions.Input{resultInput(t, turn, "a", full), {Kind: "message", Payload: json.RawMessage(`{"text":"Follow up"}`)}, resultInput(t, turn, "b", `{"success":true,"output":null,"error":null}`), {Kind: "cancel", Payload: json.RawMessage(`{}`)}} - receipts, err := s.SubmitInputs(t.Context(), tenant, session.ID, "batch", batch) + receipts, err := submitInputs(t.Context(), s, tenant, session.ID, "batch", batch) if err != nil || len(receipts) != 4 { t.Fatal(receipts, err) } @@ -60,7 +60,7 @@ func TestFunctionInputBatchesPersistAndReplayWithoutRetargeting(t *testing.T) { next := submitMessage(t, s, tenant, session.ID, "next").TurnID pool.Close() s, _ = testStore(t) - retry, err := s.SubmitInputs(t.Context(), tenant, session.ID, "batch", batch) + retry, err := submitInputs(t.Context(), s, tenant, session.ID, "batch", batch) if err != nil || len(retry) != len(receipts) { t.Fatal(retry, err) } @@ -73,11 +73,11 @@ func TestFunctionInputBatchesPersistAndReplayWithoutRetargeting(t *testing.T) { if !reflect.DeepEqual(retry, receipts) { t.Fatal(retry, receipts) } - history, err := s.ListTurnInputs(t.Context(), tenant, session.ID, turn, 0, 100) + history, err := sessionAdapter(s).ListTurnInputs(t.Context(), tenant, session.ID, turn, 0, 100) if err != nil || len(history) != 5 || history[1].Kind != "tool_result" || history[2].Kind != "message" { t.Fatal(history, err) } - future, err := s.ListTurnInputs(t.Context(), tenant, session.ID, next, 0, 100) + future, err := sessionAdapter(s).ListTurnInputs(t.Context(), tenant, session.ID, next, 0, 100) if err != nil || len(future) != 1 { t.Fatal(future, err) } @@ -86,11 +86,11 @@ func TestFunctionInputBatchesPersistAndReplayWithoutRetargeting(t *testing.T) { t.Fatal(current, err) } changed := []sessions.Input{batch[1], batch[0], batch[2], batch[3]} - if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "batch", changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { + if _, err := submitInputs(t.Context(), s, tenant, session.ID, "batch", changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal(err) } // A new request identity can repeat an identical saved result, without native application. - if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "same-result", batch[:1]); err != nil { + if _, err := submitInputs(t.Context(), s, tenant, session.ID, "same-result", batch[:1]); err != nil { t.Fatal(err) } call, err = FixtureFunctionCall(t.Context(), s.pool, tenant, session.ID, turn, "a") @@ -136,7 +136,7 @@ func TestFunctionInputBatchFailureRollsBackEveryWrite(t *testing.T) { batch = append(batch, resultInput(t, turn, "a", `{"success":false}`)) expected = sessions.ErrFunctionResultConflict } - if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "failed-batch", batch); !errors.Is(err, expected) { + if _, err := submitInputs(t.Context(), s, tenant, session.ID, "failed-batch", batch); !errors.Is(err, expected) { t.Fatal(err) } call, err := FixtureFunctionCall(t.Context(), s.pool, tenant, session.ID, turn, "a") @@ -147,11 +147,11 @@ func TestFunctionInputBatchFailureRollsBackEveryWrite(t *testing.T) { if err != nil || !state.CancelRequestedAt.IsZero() || state.Status != sessions.TurnWaiting { t.Fatal(state, err) } - history, err := s.ListTurnInputs(t.Context(), tenant, session.ID, turn, 0, 100) + history, err := sessionAdapter(s).ListTurnInputs(t.Context(), tenant, session.ID, turn, 0, 100) if err != nil || len(history) != 1 { t.Fatal(history, err) } - if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "failed-batch", []sessions.Input{first, cancel}); err != nil { + if _, err := submitInputs(t.Context(), s, tenant, session.ID, "failed-batch", []sessions.Input{first, cancel}); err != nil { t.Fatal("failed transaction retained retry identity", err) } }) @@ -169,7 +169,7 @@ func TestFunctionInputConcurrentBatchesSelectOneResult(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - _, err := other.SubmitInputs(t.Context(), tenant, session.ID, fmt.Sprint(i), batch) + _, err := submitInputs(t.Context(), other, tenant, session.ID, fmt.Sprint(i), batch) results <- err }() } @@ -188,7 +188,7 @@ func TestFunctionInputConcurrentBatchesSelectOneResult(t *testing.T) { if wins != 1 || conflicts != 1 { t.Fatal(wins, conflicts) } - history, err := s.ListTurnInputs(t.Context(), tenant, session.ID, turn, 0, 100) + history, err := sessionAdapter(s).ListTurnInputs(t.Context(), tenant, session.ID, turn, 0, 100) if err != nil || len(history) != 3 { t.Fatal(history, err) } @@ -198,7 +198,7 @@ func TestFunctionInputsRejectInvalidTargetsAndStorageObjects(t *testing.T) { s, _ := testStore(t) tenant, session, turn := functionInputFixture(t, s, functionExecution(t)) for _, raw := range []string{`{}`, `{"turn_id":"","call_id":"a","result":{}}`, fmt.Sprintf(`{"turn_id":%q,"call_id":" ","result":{}}`, turn), fmt.Sprintf(`{"turn_id":%q,"call_id":"a"}`, turn), fmt.Sprintf(`{"turn_id":%q,"call_id":"a","result":null}`, turn), fmt.Sprintf(`{"turn_id":%q,"call_id":"a","result":[]}`, turn)} { - _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "invalid", []sessions.Input{{Kind: "tool_result", Payload: json.RawMessage(raw)}}) + _, err := submitInputs(t.Context(), s, tenant, session.ID, "invalid", []sessions.Input{{Kind: "tool_result", Payload: json.RawMessage(raw)}}) if !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal(err) } @@ -207,7 +207,7 @@ func TestFunctionInputsRejectInvalidTargetsAndStorageObjects(t *testing.T) { // Missing, foreign and malformed Sessions are not found, whatever the target. for _, scope := range []struct{ tenant, session string }{{uuid.NewString(), session.ID}, {tenant, uuid.NewString()}, {tenant, "sess_malformed"}} { for _, target := range []sessions.Input{input, resultInput(t, "bad", "missing", `{"success":true}`)} { - if _, err := s.SubmitInputs(t.Context(), scope.tenant, scope.session, "foreign", []sessions.Input{target}); !errors.Is(err, sessions.ErrNotFound) { + if _, err := submitInputs(t.Context(), s, scope.tenant, scope.session, "foreign", []sessions.Input{target}); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } } @@ -221,12 +221,12 @@ func TestFunctionInputsRejectInvalidTargetsAndStorageObjects(t *testing.T) { {uuid.NewString(), "a", sessions.ErrFunctionCallTurnMismatch}, {"bad", "a", sessions.ErrFunctionCallTurnMismatch}, {turn, "missing", sessions.ErrUnknownFunctionCall}, {uuid.NewString(), "missing", sessions.ErrUnknownFunctionCall}, {"bad", "missing", sessions.ErrUnknownFunctionCall}, } { - if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "target", []sessions.Input{resultInput(t, target.turn, target.call, `{"success":true}`)}); !errors.Is(err, target.want) { + if _, err := submitInputs(t.Context(), s, tenant, session.ID, "target", []sessions.Input{resultInput(t, target.turn, target.call, `{"success":true}`)}); !errors.Is(err, target.want) { t.Fatal(target, err) } } transition(t, s, tenant, session.ID, turn, sessions.TurnWaiting, sessions.TurnFailed) - if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "late", []sessions.Input{input}); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := submitInputs(t.Context(), s, tenant, session.ID, "late", []sessions.Input{input}); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal(err) } current, err := sessionAdapter(s).GetSession(t.Context(), tenant, session.ID) diff --git a/services/core/internal/store/function_state_public_test.go b/services/core/internal/store/function_state_public_test.go index c578c5310..69b329cbf 100644 --- a/services/core/internal/store/function_state_public_test.go +++ b/services/core/internal/store/function_state_public_test.go @@ -31,7 +31,7 @@ func TestFunctionStateOfficialClientReadsAndLiveEvents(t *testing.T) { if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(ctx, tenant, session.ID, "start", json.RawMessage(`{"text":"fixture"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "start", json.RawMessage(`{"text":"fixture"}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/function_state_test.go b/services/core/internal/store/function_state_test.go index 320650e7f..8da65ee48 100644 --- a/services/core/internal/store/function_state_test.go +++ b/services/core/internal/store/function_state_test.go @@ -98,7 +98,7 @@ func TestFunctionStateCancellationAndTerminalCleanup(t *testing.T) { if status == sessions.TurnCancelled { before, _ := sessionAdapter(s).SessionEventCursor(t.Context(), tenant, session.ID) for _, key := range []string{"cancel", "cancel", "another-cancel"} { - if _, err := s.RequestCancel(t.Context(), tenant, session.ID, key); err != nil { + if _, err := requestCancel(t.Context(), s, tenant, session.ID, key); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/harness_onboarding_test.go b/services/core/internal/store/harness_onboarding_test.go index 3203b07de..c41003575 100644 --- a/services/core/internal/store/harness_onboarding_test.go +++ b/services/core/internal/store/harness_onboarding_test.go @@ -105,7 +105,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { t.Fatal(err) } var result execution.Result - inputs, inputErr := h.s.ListTurnInputs(ctx, h.tenant, created.ID, first.RunID, 0, 100) + inputs, inputErr := store.SessionAdapter(h.s).ListTurnInputs(ctx, h.tenant, created.ID, first.RunID, 0, 100) if inputErr != nil || len(inputs) != 2 { t.Fatal(inputs, inputErr) } diff --git a/services/core/internal/store/hosted_initialization_failure_public_test.go b/services/core/internal/store/hosted_initialization_failure_public_test.go index 90ccc4f5a..805d0ebb3 100644 --- a/services/core/internal/store/hosted_initialization_failure_public_test.go +++ b/services/core/internal/store/hosted_initialization_failure_public_test.go @@ -223,7 +223,7 @@ func TestHostedInitializationFailureRecordsSafeSessionFailure(t *testing.T) { !last.EnvironmentFailure.FailedAt.Equal(read.EnvironmentFailure.FailedAt) || last.EnvironmentInputActivity != nil || !last.Settled { t.Fatal("failed snapshot", last) } - if _, err := s.ReserveEnvironmentInput(t.Context(), tenant, session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}); !errors.Is(err, sessions.ErrHostedEnvironmentFailed) { + if _, err := store.SessionService(t, s).ReserveEnvironmentInput(t.Context(), tenant, session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}); !errors.Is(err, sessions.ErrHostedEnvironmentFailed) { t.Fatal("failed hosted Environment admitted input", err) } raw, _ := json.Marshal(events) diff --git a/services/core/internal/store/input_batches_test.go b/services/core/internal/store/input_batches_test.go deleted file mode 100644 index 16574d0fd..000000000 --- a/services/core/internal/store/input_batches_test.go +++ /dev/null @@ -1,189 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "reflect" - "strings" - "sync" - "testing" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" -) - -func messageInput(text string) sessions.Input { - payload, _ := json.Marshal(map[string]string{"text": text}) - return sessions.Input{Kind: "message", Payload: payload} -} - -func TestInputBatchesAreOrderedAndIdempotentAcrossConnections(t *testing.T) { - s, _ := testStore(t) - other, _ := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - const count = 8 - batch := []sessions.Input{messageInput("first"), messageInput("second")} - for _, repeated := range []bool{true, false} { - var wg sync.WaitGroup - receipts := make(chan []sessions.InputReceipt, count) - errs := make(chan error, count) - for i := range count { - wg.Add(1) - go func() { - defer wg.Done() - st := s - if i%2 == 0 { - st = other - } - key := "same-batch" - if !repeated { - key = fmt.Sprintf("batch-%d", i) - } - got, err := st.SubmitInputs(ctx, tenant, session.ID, key, batch) - receipts <- got - errs <- err - }() - } - wg.Wait() - close(receipts) - close(errs) - for err := range errs { - if err != nil { - t.Fatal(err) - } - } - newBatches := 0 - for got := range receipts { - if len(got) != 2 || got[0].TurnID == "" || got[0].TurnID != got[1].TurnID || got[0].Sequence >= got[1].Sequence || got[0].Replayed != got[1].Replayed { - t.Fatalf("invalid batch receipt: %+v", got) - } - if !got[0].Replayed { - newBatches++ - } - } - want := count - if repeated { - want = 1 - } - if newBatches != want { - t.Fatalf("accepted %d batches, want %d", newBatches, want) - } - } - got, err := s.SubmitInputs(ctx, tenant, session.ID, "same-batch", batch) - if err != nil { - t.Fatal(err) - } - inputs, err := s.ListTurnInputs(ctx, tenant, session.ID, got[0].TurnID, 0, 100) - if err != nil || len(inputs) != 2*(count+1) { - t.Fatalf("inputs=%d err=%v", len(inputs), err) - } - for i, input := range inputs { - var payload map[string]string - if err := json.Unmarshal(input.Payload, &payload); err != nil { - t.Fatal(err) - } - want := "first" - if i%2 == 1 { - want = "second" - } - if payload["text"] != want { - t.Fatalf("batch interleaved at %d: %v", i, payload) - } - } -} - -func TestBatchRetriesCompareTheWholeRequestAndRetainTargets(t *testing.T) { - s, pool := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - cancel := sessions.Input{Kind: "cancel", Payload: json.RawMessage(`{}`)} - batch := []sessions.Input{cancel, messageInput("one"), cancel, messageInput("two")} - first, err := s.SubmitInputs(ctx, tenant, session.ID, "mixed", batch) - if err != nil { - t.Fatal(err) - } - if first[0].TurnID != "" || first[1].TurnID == "" || first[1].TurnID != first[2].TurnID || first[1].TurnID == first[3].TurnID { - t.Fatalf("cancellation targets: %+v", first) - } - cancelled, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, first[1].TurnID) - if err != nil || cancelled.Status != sessions.TurnCancelled { - t.Fatalf("cancelled turn=%+v err=%v", cancelled, err) - } - transition(t, s, tenant, session.ID, first[3].TurnID, sessions.TurnQueued, sessions.TurnInProgress) - transition(t, s, tenant, session.ID, first[3].TurnID, sessions.TurnInProgress, sessions.TurnCompleted) - next := submitMessage(t, s, tenant, session.ID, "next") - for _, changed := range [][]sessions.Input{batch[:3], append(append([]sessions.Input{}, batch...), cancel), {batch[0], batch[3], batch[2], batch[1]}} { - if _, err := s.SubmitInputs(ctx, tenant, session.ID, "mixed", changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { - t.Fatalf("changed batch accepted: %v", err) - } - } - batch[1].Payload = json.RawMessage(`{ "text" : "one" }`) - pool.Close() - restarted, _ := testStore(t) - retry, err := restarted.SubmitInputs(ctx, tenant, session.ID, "mixed", batch) - for i := range first { - first[i].Replayed = true - } - if err != nil || !reflect.DeepEqual(retry, first) { - t.Fatalf("restart changed receipts: %+v, %v", retry, err) - } - current, err := sessionAdapter(restarted).GetTurn(ctx, tenant, session.ID, next.TurnID) - if err != nil || current.Status != sessions.TurnQueued || !current.CancelRequestedAt.IsZero() { - t.Fatalf("retry cancelled later work: %+v, %v", current, err) - } - otherTenant, _ := newTurnSession(t, restarted) - if _, err := restarted.SubmitInputs(ctx, otherTenant, session.ID, "mixed", batch); !errors.Is(err, sessions.ErrNotFound) { - t.Fatalf("batch retry escaped tenant: %v", err) - } -} - -func TestFailedBatchRollsBackEarlierCancellationAndInputs(t *testing.T) { - s, pool := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - initial := submitMessage(t, s, tenant, session.ID, "initial") - transition(t, s, tenant, session.ID, initial.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - key := uuid.NewString() - constraint := "batch_failure_" + strings.ReplaceAll(key, "-", "") - // Inject a storage error on the second insert, after cancellation has run. - _, err := pool.Exec(ctx, "ALTER TABLE turn_inputs ADD CONSTRAINT "+constraint+" CHECK (idempotency_key <> '"+key+"' OR batch_position = 0)") - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _, _ = pool.Exec(ctx, "ALTER TABLE turn_inputs DROP CONSTRAINT IF EXISTS "+constraint) }) - batch := []sessions.Input{{Kind: "cancel", Payload: json.RawMessage(`{}`)}, messageInput("after cancel")} - if got, err := s.SubmitInputs(ctx, tenant, session.ID, key, batch); err == nil || got != nil { - t.Fatalf("partial batch succeeded: %+v %v", got, err) - } - turn, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, initial.TurnID) - if err != nil || turn.Status != sessions.TurnInProgress || !turn.CancelRequestedAt.IsZero() { - t.Fatalf("cancellation escaped rollback: %+v %v", turn, err) - } - inputs, err := s.ListTurnInputs(ctx, tenant, session.ID, initial.TurnID, 0, 100) - if err != nil || len(inputs) != 1 { - t.Fatalf("partial inputs survived: %v %v", inputs, err) - } - if _, err := pool.Exec(ctx, "ALTER TABLE turn_inputs DROP CONSTRAINT "+constraint); err != nil { - t.Fatal(err) - } - got, err := s.SubmitInputs(ctx, tenant, session.ID, key, batch) - if err != nil || len(got) != 2 || got[0].Replayed || got[1].Replayed { - t.Fatalf("retry after rollback: %+v %v", got, err) - } -} - -func TestInputBatchValidation(t *testing.T) { - for _, input := range [][]sessions.Input{ - nil, make([]sessions.Input, 65), {messageInput("ok"), {Kind: "unsupported", Payload: json.RawMessage(`{}`)}}, - {{Kind: "message"}}, {{Kind: "message", Payload: json.RawMessage(`[]`)}}, - {{Kind: "cancel", Payload: json.RawMessage(`{"target":"other"}`)}}, - {messageInput(strings.Repeat("x", 300*1024)), messageInput(strings.Repeat("y", 300*1024))}, - } { - if _, _, err := validateInputs(input); !errors.Is(err, sessions.ErrInvalidInput) { - t.Fatalf("invalid batch accepted: %v", err) - } - } -} diff --git a/services/core/internal/store/input_conflicts_public_test.go b/services/core/internal/store/input_conflicts_public_test.go index fbe1e422c..7c739e4b1 100644 --- a/services/core/internal/store/input_conflicts_public_test.go +++ b/services/core/internal/store/input_conflicts_public_test.go @@ -70,7 +70,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { // waiting starts a Turn that waits for one function result. waiting := func(session, key, call string) sessions.InputReceipt { t.Helper() - receipt, err := s.SubmitMessage(ctx, tenant, session, key, json.RawMessage(`{"text":"work"}`)) + receipt, err := store.SendMessage(ctx, s, tenant, session, key, json.RawMessage(`{"text":"work"}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/inputs_test.go b/services/core/internal/store/inputs_test.go new file mode 100644 index 000000000..ec00d3d78 --- /dev/null +++ b/services/core/internal/store/inputs_test.go @@ -0,0 +1,128 @@ +package store + +import ( + "context" + "encoding/json" + "testing" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var messagePayload = json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"hello"}]}]}`) + +func messageInput(text string) sessions.Input { + payload, _ := json.Marshal(map[string]string{"text": text}) + return sessions.Input{Kind: "message", Payload: payload} +} + +func newTurnSession(t *testing.T, s *Store) (string, sessions.Session) { + t.Helper() + tenant := uuid.NewString() + session, err := s.CreateSession(context.Background(), tenant, sessions.CreateSession{Creator: FixtureCreator(), + Engine: "codex", IdempotencyKey: "session", Configuration: json.RawMessage(`{"agent":{"model":"test","instructions":"original"}}`), + }) + if err != nil { + t.Fatal(err) + } + return tenant, session +} + +func environmentInputSession(t *testing.T, s *Store) (string, sessions.Session) { + t.Helper() + tenant := uuid.NewString() + session, err := s.CreateSession(context.Background(), tenant, environmentInput("session", "self_hosted", "/workspace")) + if err != nil { + t.Fatal(err) + } + return tenant, session +} + +// submitInputs admits inputs through the Session service on s's database, as +// the public events route does. +func submitInputs(ctx context.Context, s *Store, tenant, session, key string, inputs []sessions.Input) ([]sessions.InputReceipt, error) { + service, err := sessions.NewService(sessionAdapter(s)) + if err != nil { + return nil, err + } + return service.SubmitInputs(ctx, tenant, session, key, inputs) +} + +func submitInput(ctx context.Context, s *Store, tenant, session, key string, input sessions.Input) (sessions.InputReceipt, error) { + receipts, err := submitInputs(ctx, s, tenant, session, key, []sessions.Input{input}) + if err != nil { + return sessions.InputReceipt{}, err + } + return receipts[0], nil +} + +func sendMessage(ctx context.Context, s *Store, tenant, session, key string, payload json.RawMessage) (sessions.InputReceipt, error) { + return submitInput(ctx, s, tenant, session, key, sessions.Input{Kind: "message", Payload: payload}) +} + +func requestCancel(ctx context.Context, s *Store, tenant, session, key string) (sessions.InputReceipt, error) { + return submitInput(ctx, s, tenant, session, key, sessions.Input{Kind: "cancel", Payload: json.RawMessage(`{}`)}) +} + +func submitMessage(t *testing.T, s *Store, tenant, session, key string) sessions.InputReceipt { + t.Helper() + receipt, err := sendMessage(context.Background(), s, tenant, session, key, messagePayload) + if err != nil { + t.Fatal(err) + } + return receipt +} + +func reserveEnvironmentInput(t *testing.T, s *Store, tenant, session, key string) sessions.EnvironmentInputReservation { + t.Helper() + got, err := sessionService(t, s).ReserveEnvironmentInput(context.Background(), tenant, session, key, []sessions.Input{messageInput("first"), messageInput("second")}) + if err != nil { + t.Fatal(err) + } + return got +} + +// cancelEnvironmentInput cancels the Session's pending Environment input as +// Session cancellation does, then reads the reservation back. +func cancelEnvironmentInput(ctx context.Context, s *Store, tenant, session, reservation string) (sessions.EnvironmentInputReservation, error) { + owner, err := parseID(tenant) + if err != nil { + return sessions.EnvironmentInputReservation{}, err + } + err = s.withPublicSession(ctx, tenant, session, func(ctx context.Context, q *sqlc.Queries, id pgtype.UUID) error { + bound := sessionpg.BindSession(q, owner, id) + return sessions.TrackInputActivity(ctx, bound, func(ctx context.Context) error { return bound.CancelPendingInput(ctx) }) + }) + if err != nil { + return sessions.EnvironmentInputReservation{}, err + } + return sessionAdapter(s).GetEnvironmentInputReservation(ctx, tenant, session, reservation) +} + +func transition(t *testing.T, s *Store, tenant, session, turn, from, to string) sessions.Turn { + t.Helper() + got, err := transitionTurn(context.Background(), s, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: from, Status: to}) + if err != nil { + t.Fatal(err) + } + return got +} + +func environmentInputHistory(t *testing.T, pool *pgxpool.Pool, session string, turns, inputs int) { + t.Helper() + var gotTurns, gotInputs, items, events int + err := pool.QueryRow(context.Background(), ` + SELECT (SELECT count(*) FROM turns WHERE session_id=$1), + (SELECT count(*) FROM turn_inputs WHERE session_id=$1), + (SELECT count(*) FROM session_items WHERE session_id=$1), + (SELECT count(*) FROM session_events WHERE session_id=$1 + AND (payload ? 'turn' OR payload->'event' ? 'item'))`, session).Scan(&gotTurns, &gotInputs, &items, &events) + if err != nil || gotTurns != turns || gotInputs != inputs || items != inputs || (inputs == 0 && events != 0) { + t.Fatal("history", gotTurns, gotInputs, items, events, err) + } +} diff --git a/services/core/internal/store/item_order_test.go b/services/core/internal/store/item_order_test.go index 924da4b50..5514b2e31 100644 --- a/services/core/internal/store/item_order_test.go +++ b/services/core/internal/store/item_order_test.go @@ -23,7 +23,7 @@ func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(ctx, tenant, session.ID, "first", json.RawMessage(`{"text":"question"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "first", json.RawMessage(`{"text":"question"}`)) if err != nil { t.Fatal(err) } @@ -52,7 +52,7 @@ func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) t.Fatal(err) } } - if _, err = s.SubmitMessage(ctx, tenant, session.ID, "steer", json.RawMessage(`{"text":"continue"}`)); err != nil { + if _, err = store.SendMessage(ctx, s, tenant, session.ID, "steer", json.RawMessage(`{"text":"continue"}`)); err != nil { t.Fatal(err) } page, err = sessionReads(pool).ListItems(ctx, tenant, session.ID, "", 100, true) @@ -132,7 +132,7 @@ func TestItemObservationOrderSurvivesTiesUpdatesRetriesAndRecovery(t *testing.T) t.Fatal(err) } checkOrder() - next, err := s.SubmitMessage(ctx, tenant, session.ID, "next-turn", json.RawMessage(`{"text":"new turn"}`)) + next, err := store.SendMessage(ctx, s, tenant, session.ID, "next-turn", json.RawMessage(`{"text":"new turn"}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/item_reads_test.go b/services/core/internal/store/item_reads_test.go index 836fd4333..861f0d7f4 100644 --- a/services/core/internal/store/item_reads_test.go +++ b/services/core/internal/store/item_reads_test.go @@ -23,7 +23,7 @@ func TestItemsRecoverSnapshotsPartialResultsPaginationAndIsolation(t *testing.T) if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(ctx, tenant, session.ID, "first", json.RawMessage(`{"text":"question"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "first", json.RawMessage(`{"text":"question"}`)) if err != nil { t.Fatal(err) } @@ -123,7 +123,7 @@ func TestItemProjectionFailureRollsBackJournalAndAggregateRecovers(t *testing.T) journal := executionOwner(t, fixtureDB{pool: pool}, s).Sessions tenant := uuid.NewString() session, _ := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: "legacy"}) - input, err := s.SubmitMessage(ctx, tenant, session.ID, "input", json.RawMessage(`{"text":"test"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "input", json.RawMessage(`{"text":"test"}`)) if err != nil { t.Fatal(err) } @@ -160,7 +160,7 @@ func TestReceiptOnlyTextRecoversWithoutInventingCompletion(t *testing.T) { tenant := uuid.NewString() for _, receiptOnly := range []bool{true, false} { session, _ := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: store.FixtureCreator(), Engine: "codex", IdempotencyKey: uuid.NewString()}) - input, err := s.SubmitMessage(ctx, tenant, session.ID, "first", json.RawMessage(`{"text":"test"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "first", json.RawMessage(`{"text":"test"}`)) if err != nil { t.Fatal(err) } @@ -194,7 +194,7 @@ func TestLegacyFailureRetainsPartialAnswerAcrossRecovery(t *testing.T) { if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(ctx, tenant, session.ID, "first", json.RawMessage(`{"text":"question"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "first", json.RawMessage(`{"text":"question"}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/list_cursor_public_test.go b/services/core/internal/store/list_cursor_public_test.go index 70d58d2dc..8d050877c 100644 --- a/services/core/internal/store/list_cursor_public_test.go +++ b/services/core/internal/store/list_cursor_public_test.go @@ -69,7 +69,7 @@ func seedCursorFixture(t *testing.T, s *store.Store, leased execution.Owner, ski if _, err := store.TransitionTurn(ctx, s, tenant, f.session, f.turn, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnCancelled}); err != nil { t.Fatal(err) } - if _, err := s.SubmitMessage(ctx, tenant, f.session, label+"-second", json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"second"}]}]}`)); err != nil { + if _, err := store.SendMessage(ctx, s, tenant, f.session, label+"-second", json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"second"}]}]}`)); err != nil { t.Fatal(err) } f.otherSession = client.created(token, "/v1/agents/sessions", newSession) @@ -91,7 +91,7 @@ func seedCursorFixture(t *testing.T, s *store.Store, leased execution.Owner, ski if err != nil { t.Fatal(err) } - receipt, err := s.SubmitMessage(ctx, tenant, created.ID, key+"-input", json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"delegate"}]}]}`)) + receipt, err := store.SendMessage(ctx, s, tenant, created.ID, key+"-input", json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"delegate"}]}]}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/local_artifact_export_test.go b/services/core/internal/store/local_artifact_export_test.go index cc45e09e5..b39f29ded 100644 --- a/services/core/internal/store/local_artifact_export_test.go +++ b/services/core/internal/store/local_artifact_export_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" ) func completeLocalArtifactExport(t *testing.T, h *dispatchHarness, worker *execution.Worker, environment sessions.Environment) { @@ -46,7 +47,7 @@ func completeLocalArtifactExport(t *testing.T, h *dispatchHarness, worker *execu t.Fatal("capture published before native completion", err) } completeCaptureDirectoryRead(t, h, worker, environment) - pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "during-artifact-capture", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"run after the completed native execution"}`)}}) + pending, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "during-artifact-capture", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"run after the completed native execution"}`)}}) if err != nil || pending.State != sessions.EnvironmentInputPending || len(pending.Receipts) != 0 { t.Fatalf("input during artifact capture was assigned to the finished executor: %+v %v", pending, err) } 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 ed77c1edf..dd795c914 100644 --- a/services/core/internal/store/local_environment_file_write_test.go +++ b/services/core/internal/store/local_environment_file_write_test.go @@ -57,7 +57,7 @@ func TestLocalEnvironmentFileWriteOwnsMutationBeforeDispatch(t *testing.T) { if err != nil || intent.State != "pending" || intent.Identity.DeviceID != h.device.ID { t.Fatal("dispatch preceded durable ownership", intent, err) } - if _, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "concurrent", []sessions.Input{{Kind: "message", Payload: []byte(`{"text":"work"}`)}}); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "concurrent", []sessions.Input{{Kind: "message", Payload: []byte(`{"text":"work"}`)}}); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("upload admitted concurrent execution", err) } cancel() @@ -99,7 +99,7 @@ func TestLocalEnvironmentFileWriteLostReceiptRemainsPending(t *testing.T) { if err != nil || intent.State != "pending" { t.Fatal("disconnect guessed rejection", intent, err) } - if _, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "after-loss", []sessions.Input{{Kind: "message", Payload: []byte(`{"text":"work"}`)}}); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "after-loss", []sessions.Input{{Kind: "message", Payload: []byte(`{"text":"work"}`)}}); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("unknown upload admitted execution", err) } } diff --git a/services/core/internal/store/local_environment_worker_test.go b/services/core/internal/store/local_environment_worker_test.go index cd8e2429a..6f935956d 100644 --- a/services/core/internal/store/local_environment_worker_test.go +++ b/services/core/internal/store/local_environment_worker_test.go @@ -110,7 +110,7 @@ func TestLocalEnvironmentWorkerRejectsGeneralDeviceDespiteCapability(t *testing. func TestLocalEnvironmentWorkerSchedulesPreparationWithoutRemoteResolver(t *testing.T) { h, worker, environment := localWorker(t, true, true) - reservation, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "local-input", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) + reservation, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "local-input", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) if err != nil { t.Fatal(err) } @@ -144,7 +144,7 @@ func TestLocalEnvironmentWorkerSchedulesPreparationWithoutRemoteResolver(t *test turn, err := store.SessionAdapter(h.s).GetTurn(t.Context(), h.tenant, h.session.ID, start.RunID) return err == nil && turn.Status == sessions.TurnCompleted }) - settled, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, reservation.ID) + settled, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, reservation.ID) if err != nil || settled.State != sessions.EnvironmentInputAdmitted || len(settled.Receipts) != 1 { t.Fatal("local reservation did not settle", err) } diff --git a/services/core/internal/store/mcode_public_native_test.go b/services/core/internal/store/mcode_public_native_test.go index 9cd0dcd28..9e1b3ce01 100644 --- a/services/core/internal/store/mcode_public_native_test.go +++ b/services/core/internal/store/mcode_public_native_test.go @@ -89,7 +89,7 @@ func TestNativeMCodePublicExecution(t *testing.T) { if err != nil { t.Fatal(err) } - inputs, err := h.s.ListTurnInputs(ctx, h.tenant, proof.Session, proof.FirstTurn, 0, 100) + inputs, err := store.SessionAdapter(h.s).ListTurnInputs(ctx, h.tenant, proof.Session, proof.FirstTurn, 0, 100) if err != nil || len(inputs) != 2 { t.Fatal("steering input not in same turn", err) } diff --git a/services/core/internal/store/message_images_native_test.go b/services/core/internal/store/message_images_native_test.go index 40019ebc2..79aa40647 100644 --- a/services/core/internal/store/message_images_native_test.go +++ b/services/core/internal/store/message_images_native_test.go @@ -86,7 +86,7 @@ func TestNativeMessageImagePublicExecution(t *testing.T) { if json.Unmarshal(turn.Outcome, &outcome) != nil || outcome.AppliedThrough < 1 { t.Fatal("native input receipt missing") } - inputs, err := h.s.ListTurnInputs(ctx, h.tenant, proof.Session, proof.Turn, 0, 100) + inputs, err := store.SessionAdapter(h.s).ListTurnInputs(ctx, h.tenant, proof.Session, proof.Turn, 0, 100) if err != nil || len(inputs) != 3 || inputs[0].Kind != "message" || inputs[1].Kind != "message" || inputs[2].Kind != "tool_result" || outcome.AppliedThrough != inputs[2].Sequence { t.Fatal("active image batch was not applied exactly once in the same Turn", err) } diff --git a/services/core/internal/store/model_protocol_native_test.go b/services/core/internal/store/model_protocol_native_test.go index d0ee2dbd2..fac467b0e 100644 --- a/services/core/internal/store/model_protocol_native_test.go +++ b/services/core/internal/store/model_protocol_native_test.go @@ -125,7 +125,7 @@ func TestNativeModelProtocolPublicExecution(t *testing.T) { if err != nil || !call.Applied { t.Fatal("public function result lacks native delivery acknowledgement") } - inputs, err := h.s.ListTurnInputs(ctx, h.tenant, proof.Session, item.Turn, 0, 100) + inputs, err := store.SessionAdapter(h.s).ListTurnInputs(ctx, h.tenant, proof.Session, item.Turn, 0, 100) if err != nil || len(inputs) != 2 || inputs[0].Kind != "message" || inputs[1].Kind != "tool_result" { t.Fatal("public function result input was lost or duplicated") } diff --git a/services/core/internal/store/prepared_dispatch_failure_test.go b/services/core/internal/store/prepared_dispatch_failure_test.go index aef19d944..4c62c3518 100644 --- a/services/core/internal/store/prepared_dispatch_failure_test.go +++ b/services/core/internal/store/prepared_dispatch_failure_test.go @@ -24,7 +24,7 @@ func TestPreparedDispatchSettlesOnlyReadyInput(t *testing.T) { handle := acknowledgePreparation(h, frame.ID) switch action { case "cancel": - if _, err := h.s.CancelEnvironmentInput(t.Context(), h.tenant, h.session.ID, pending.ID); err != nil { + if _, err := store.CancelEnvironmentInput(t.Context(), h.s, h.tenant, h.session.ID, pending.ID); err != nil { t.Fatal(err) } case "expire": @@ -66,7 +66,7 @@ func TestPreparedDispatchSettlesOnlyReadyInput(t *testing.T) { } assertEnvironmentExpiryHasNoHistory(t, pool, h.session.ID) if action == "prepare-failure" || action == "disconnect" { - stored, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) + stored, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) if err != nil || stored.State != sessions.EnvironmentInputPending || !stored.Deadline.Equal(pending.Deadline) { t.Fatal("preparation failure changed the pending identity", stored, err) } @@ -87,7 +87,7 @@ func TestPreparedDispatchHandlesStartRejectionAndPendingStartCancellation(t *tes h.write(frame.ID, proto.TypePreparationStatus, proto.PreparationStatusPayload{State: "rejected", Operation: proto.TypeExecutionStart, ErrorCode: "preparation_not_ready"}) } else { h.write(frame.ID, proto.TypePreparationStatus, proto.PreparationStatusPayload{Handle: handle, Revision: 3, State: "starting", RunID: start.RunID}) - if _, err := h.s.RequestCancel(t.Context(), h.tenant, h.session.ID, "cancel-start"); err != nil { + if _, err := store.RequestCancel(t.Context(), h.s, h.tenant, h.session.ID, "cancel-start"); err != nil { t.Fatal(err) } frame := h.read(proto.TypePromptCancel) @@ -147,7 +147,7 @@ func TestPreparedDispatchCancellationReceiptSurvivesStartFailure(t *testing.T) { handle := acknowledgePreparation(h, prepare.ID) start := readyPreparedDispatch(t, h, prepare.ID, handle) h.write(prepare.ID, proto.TypePreparationStatus, proto.PreparationStatusPayload{Handle: handle, Revision: 3, State: "starting", RunID: start.RunID}) - if _, err := h.s.RequestCancel(t.Context(), h.tenant, h.session.ID, "cancel-start"); err != nil { + if _, err := store.RequestCancel(t.Context(), h.s, h.tenant, h.session.ID, "cancel-start"); err != nil { t.Fatal(err) } frame := h.read(proto.TypePromptCancel) diff --git a/services/core/internal/store/prepared_dispatch_test.go b/services/core/internal/store/prepared_dispatch_test.go index 4c30eadc6..5964853c0 100644 --- a/services/core/internal/store/prepared_dispatch_test.go +++ b/services/core/internal/store/prepared_dispatch_test.go @@ -25,7 +25,7 @@ func preparedDispatchHarness(t *testing.T) (*dispatchHarness, sessions.Environme assertNoRuntimeAllocation(t, h) h.d, h.lease = h.bound(), h.owner().Lease enableWorkerEnvironment(t, h) - pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}, {Kind: "message", Payload: json.RawMessage(`{"text":"second"}`)}}) + pending, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}, {Kind: "message", Payload: json.RawMessage(`{"text":"second"}`)}}) if err != nil { t.Fatal(err) } @@ -128,7 +128,7 @@ func TestPreparedDispatchOwnerOutlivesReservationDeadline(t *testing.T) { if _, err := pool.Exec(t.Context(), "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE id=$1", pending.ID); err != nil { t.Fatal(err) } - stored, err := h.d.Store.ExpireEnvironmentInput(t.Context(), h.tenant, h.session.ID, pending.ID) + stored, err := h.d.Sessions.ExpireEnvironmentInput(t.Context(), h.tenant, h.session.ID, pending.ID) if err != nil || stored.State != sessions.EnvironmentInputAdmitted { t.Fatal("admitted execution lost its owner to the pending-input deadline", err) } diff --git a/services/core/internal/store/public_execution_test.go b/services/core/internal/store/public_execution_test.go index 6b6d99340..681841b58 100644 --- a/services/core/internal/store/public_execution_test.go +++ b/services/core/internal/store/public_execution_test.go @@ -145,7 +145,7 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { } checkMeasurement(false) queued := publicSession(t, h, "queued") - if _, err := h.s.SubmitMessage(ctx, h.tenant, queued.ID, "first", json.RawMessage(`{"text":"Not sent"}`)); err != nil { + if _, err := store.SendMessage(ctx, h.s, h.tenant, queued.ID, "first", json.RawMessage(`{"text":"Not sent"}`)); err != nil { t.Fatal(err) } worker := startOwnedWorker(t, ctx, h.db, h.d, h.owner()) @@ -162,7 +162,7 @@ func TestWorkerRestartReconcilesClaimedButPreservesQueuedWork(t *testing.T) { if err != nil || pending.LastTurn.Status != sessions.TurnQueued { t.Fatal(pending, err) } - if _, err := h.s.RequestCancel(ctx, h.tenant, queued.ID, "stop-before-dispatch"); err != nil { + if _, err := store.RequestCancel(ctx, h.s, h.tenant, queued.ID, "stop-before-dispatch"); err != nil { t.Fatal(err) } pending, err = store.SessionAdapter(h.s).GetSession(ctx, h.tenant, queued.ID) diff --git a/services/core/internal/store/public_handler_fixture_test.go b/services/core/internal/store/public_handler_fixture_test.go index 198938cbd..594490d71 100644 --- a/services/core/internal/store/public_handler_fixture_test.go +++ b/services/core/internal/store/public_handler_fixture_test.go @@ -145,14 +145,14 @@ func withPolicy(policy execution.Policy) func(*api.Dependencies) { return func(d *api.Dependencies) { d.Policy = policy } } -// storeExecution admits Sessions and inputs through the Store without a -// Worker, so nothing runs them. +// storeExecution admits Sessions through the Store and inputs through the +// Session service without a Worker, so nothing runs them. func storeExecution(t testing.TB, s *store.Store) func(*api.Dependencies) { return func(d *api.Dependencies) { d.Execution = &api.Execution{ ExecutorURL: testExecutorURL, SessionAdmission: s, - InputAdmission: s, + InputAdmission: store.SessionService(t, s), SessionArchive: strictStandIn{t}, Workspaces: strictStandIn{t}, } diff --git a/services/core/internal/store/runtime_environment_terminal_test.go b/services/core/internal/store/runtime_environment_terminal_test.go index 3a1ed2d7f..d90215a2d 100644 --- a/services/core/internal/store/runtime_environment_terminal_test.go +++ b/services/core/internal/store/runtime_environment_terminal_test.go @@ -53,11 +53,11 @@ func TestManagedEnvironmentTerminationSettlesInputAndPreservesIdentity(t *testin if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), owner.DeviceID); err != nil || ok { t.Fatal("terminal credential remained usable", err) } - failed, err := writer.PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) + failed, err := sessionExecution(t, writer.lease).PromoteEnvironmentInput(t.Context(), tenant, session.ID, reservation.ID) if err != nil || failed.State != sessions.EnvironmentInputFailed || len(failed.Receipts) != 0 || failed.SettledAt == nil || !failed.Deadline.Equal(reservation.Deadline) { t.Fatal("late preparation resurrected failed input", failed, err) } - if _, err := s.ReserveEnvironmentInput(t.Context(), tenant, session.ID, "new", []sessions.Input{messageInput("later")}); !errors.Is(err, sessions.ErrEnvironmentUnavailable) || expired == errors.Is(err, sessions.ErrHostedEnvironmentFailed) { + if _, err := sessionService(t, s).ReserveEnvironmentInput(t.Context(), tenant, session.ID, "new", []sessions.Input{messageInput("later")}); !errors.Is(err, sessions.ErrEnvironmentUnavailable) || expired == errors.Is(err, sessions.ErrHostedEnvironmentFailed) { t.Fatal("terminal environment admitted new input", err) } if _, err := s.CreateSession(t.Context(), tenant, input); err != nil { diff --git a/services/core/internal/store/runtime_input_admission_test.go b/services/core/internal/store/runtime_input_admission_test.go index b80696f0a..346d3c3c2 100644 --- a/services/core/internal/store/runtime_input_admission_test.go +++ b/services/core/internal/store/runtime_input_admission_test.go @@ -17,7 +17,7 @@ func TestManagedRuntimeMaintenancePreservesCancelAndRetry(t *testing.T) { s, db := newManagedTestStoreDB(t) tenant, session, _ := managedSession(t, s, db) inputs := []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"accepted work"}`)}} - accepted, err := s.SubmitInputs(t.Context(), tenant, session.ID, "work", inputs) + accepted, err := store.SubmitInputs(t.Context(), s, tenant, session.ID, "work", inputs) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/runtime_suspension_concurrency_test.go b/services/core/internal/store/runtime_suspension_concurrency_test.go index 67aaf940f..0bd57d443 100644 --- a/services/core/internal/store/runtime_suspension_concurrency_test.go +++ b/services/core/internal/store/runtime_suspension_concurrency_test.go @@ -99,16 +99,16 @@ func TestRuntimeSuspensionPromotionRetainsPendingInput(t *testing.T) { s, w, pool, owner := runtimeSuspensionFixture(t) runtimeSuspensionSQL(t, pool, `UPDATE runtime_allocations SET compute_phase='waking',compute_retained_until=clock_timestamp()+interval '1 hour' WHERE id=$1`, owner.ID) pending := reserveEnvironmentInput(t, s, owner.TenantID, owner.SessionID, "during-wake") - if _, err := w.PromoteEnvironmentInput(t.Context(), owner.TenantID, owner.SessionID, pending.ID); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := sessionExecution(t, w.lease).PromoteEnvironmentInput(t.Context(), owner.TenantID, owner.SessionID, pending.ID); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal("waking allocation promoted model input", err) } - got, err := s.GetEnvironmentInputReservation(t.Context(), owner.TenantID, owner.SessionID, pending.ID) + got, err := sessionAdapter(s).GetEnvironmentInputReservation(t.Context(), owner.TenantID, owner.SessionID, pending.ID) if err != nil || got.State != sessions.EnvironmentInputPending || got.SettledAt != nil || len(got.Receipts) != 0 { t.Fatal("blocked promotion partially committed", got, err) } environmentInputHistory(t, pool, owner.SessionID, 0, 0) runtimeSuspensionSQL(t, pool, `UPDATE runtime_allocations SET compute_phase='running',compute_retained_until=NULL WHERE id=$1`, owner.ID) - got, err = w.PromoteEnvironmentInput(t.Context(), owner.TenantID, owner.SessionID, pending.ID) + got, err = sessionExecution(t, w.lease).PromoteEnvironmentInput(t.Context(), owner.TenantID, owner.SessionID, pending.ID) if err != nil || got.State != sessions.EnvironmentInputAdmitted || len(got.Receipts) != 2 { t.Fatal("pending request could not resume once running", got, err) } diff --git a/services/core/internal/store/runtime_wake_hint_integration_test.go b/services/core/internal/store/runtime_wake_hint_integration_test.go index d6744f12d..b7cfc0664 100644 --- a/services/core/internal/store/runtime_wake_hint_integration_test.go +++ b/services/core/internal/store/runtime_wake_hint_integration_test.go @@ -13,6 +13,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) type wakeHintScanProvider struct { @@ -125,7 +126,7 @@ func (f *wakeHintIntegration) pending(t *testing.T, target wakeHintIntegrationTa "SELECT id::text FROM environment_input_reservations WHERE session_id=$1 AND idempotency_key=$2", target.session.ID, key).Scan(&id) == nil }) - pending, err := f.fixture.store.GetEnvironmentInputReservation(t.Context(), target.tenant, target.session.ID, id) + pending, err := store.SessionAdapter(f.fixture.store).GetEnvironmentInputReservation(t.Context(), target.tenant, target.session.ID, id) if err != nil || pending.State != sessions.EnvironmentInputPending { t.Fatal("input was not durably pending", pending.State, err) } @@ -196,7 +197,7 @@ func TestRuntimeWakeHintRejectedSubmitDoesNotAccelerateScan(t *testing.T) { } else { // Persist directly while the sentinel is blocked. Only the failing // Worker submission could emit a hint; Store persistence cannot. - if _, err := f.fixture.store.ReserveEnvironmentInput(t.Context(), f.target.tenant, f.target.session.ID, key, inputs); err != nil { + if _, err := store.SessionService(t, f.fixture.store).ReserveEnvironmentInput(t.Context(), f.target.tenant, f.target.session.ID, key, inputs); err != nil { t.Fatal(err) } if name == "idempotency conflict" { diff --git a/services/core/internal/store/runtime_worker_recovery_test.go b/services/core/internal/store/runtime_worker_recovery_test.go index 50c70e6fe..99e4a3275 100644 --- a/services/core/internal/store/runtime_worker_recovery_test.go +++ b/services/core/internal/store/runtime_worker_recovery_test.go @@ -36,7 +36,7 @@ func runtimeWorkerHarness(t *testing.T) (*dispatchHarness, *pgxpool.Pool) { func TestPreparedDispatchKeepsPendingReservationAfterComputeConflict(t *testing.T) { h, _ := runtimeWorkerHarness(t) h.d, h.lease = h.bound(), h.owner().Lease - pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) + pending, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) if err != nil { t.Fatal(err) } @@ -57,7 +57,7 @@ func TestPreparedDispatchKeepsPendingReservationAfterComputeConflict(t *testing. func TestWorkerWaitsForComputeAndSurvivesPromotionConflict(t *testing.T) { h, pool := runtimeWorkerHarness(t) - pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) + pending, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) if err != nil { t.Fatal(err) } @@ -104,7 +104,7 @@ func TestWorkerWaitsForComputeAndSurvivesPromotionConflict(t *testing.T) { if release.ID != prepare.ID { t.Fatal("conflicted preparation was not released") } - stored, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) + stored, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) if err != nil || stored.State != sessions.EnvironmentInputPending || len(stored.Receipts) != 0 { t.Fatal("conflict consumed queued input", stored, err) } diff --git a/services/core/internal/store/sandbox_deployment_switch_test.go b/services/core/internal/store/sandbox_deployment_switch_test.go index 2c26d8739..63602a7ed 100644 --- a/services/core/internal/store/sandbox_deployment_switch_test.go +++ b/services/core/internal/store/sandbox_deployment_switch_test.go @@ -311,7 +311,7 @@ func TestSandboxSwitchPreservesReleasedAllocationAndItemHistory(t *testing.T) { if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(t.Context(), tenant, history.ID, "history", json.RawMessage(`{"text":"retained request"}`)) + input, err := sendMessage(t.Context(), s, tenant, history.ID, "history", json.RawMessage(`{"text":"retained request"}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/scheduling.go b/services/core/internal/store/scheduling.go deleted file mode 100644 index b094ed530..000000000 --- a/services/core/internal/store/scheduling.go +++ /dev/null @@ -1,26 +0,0 @@ -package store - -import ( - "context" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" -) - -func (s *Store) ListEnvironmentInputWork(ctx context.Context, after string, connectedDevices []string) ([]sessions.EnvironmentInputWork, error) { - id, devices, err := sessionpg.ExecutionWorkCursor(after, connectedDevices) - if err != nil { - return nil, err - } - rows, err := s.queries.ListEnvironmentInputWork(ctx, sqlc.ListEnvironmentInputWorkParams{AfterID: id, ConnectedDevices: devices}) - if err != nil { - return nil, err - } - work := make([]sessions.EnvironmentInputWork, 0, len(rows)) - for _, row := range rows { - work = append(work, sessions.EnvironmentInputWork{TenantID: uuid.UUID(row.TenantID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), ReservationID: uuid.UUID(row.ID.Bytes).String()}) - } - return work, nil -} diff --git a/services/core/internal/store/self_hosted_cancel_public_test.go b/services/core/internal/store/self_hosted_cancel_public_test.go index 5da4bc866..aeff47ca7 100644 --- a/services/core/internal/store/self_hosted_cancel_public_test.go +++ b/services/core/internal/store/self_hosted_cancel_public_test.go @@ -118,7 +118,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { return value } idleReceipts := receipts(created.IdleKey, "") - later, err := s.ReserveEnvironmentInput(t.Context(), tenant, created.LaterID, "controlled-later-input", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"Retain pending input."}`)}}) + later, err := store.SessionService(t, s).ReserveEnvironmentInput(t.Context(), tenant, created.LaterID, "controlled-later-input", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"Retain pending input."}`)}}) if err != nil || later.State != sessions.EnvironmentInputPending || later.IsInitial { t.Fatal("could not establish controlled later reservation", err) } @@ -132,7 +132,7 @@ func TestSelfHostedCancellationOfficialClient(t *testing.T) { start := func() string { t.Helper() // Controlled callbacks isolate HTTP admission; no daemon or model runs in this fixture. - input, err := s.SubmitMessage(t.Context(), tenant, created.ID, uuid.NewString(), json.RawMessage(`{"text":"Controlled active work."}`)) + input, err := store.SendMessage(t.Context(), s, tenant, created.ID, uuid.NewString(), json.RawMessage(`{"text":"Controlled active work."}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/self_hosted_initial_public_test.go b/services/core/internal/store/self_hosted_initial_public_test.go index 174f9d687..ef1b5e050 100644 --- a/services/core/internal/store/self_hosted_initial_public_test.go +++ b/services/core/internal/store/self_hosted_initial_public_test.go @@ -107,7 +107,7 @@ func TestSelfHostedInitialCreationOfficialClient(t *testing.T) { if err := pool.QueryRow(t.Context(), "SELECT id FROM environment_input_reservations WHERE session_id=$1 AND is_initial", item.ID).Scan(&id); err != nil { t.Fatal("public creation did not reserve initial input", err) } - reservation, err := s.GetEnvironmentInputReservation(t.Context(), tenant, item.ID, id) + reservation, err := store.SessionAdapter(s).GetEnvironmentInputReservation(t.Context(), tenant, item.ID, id) if err != nil || !reservation.IsInitial || len(reservation.Inputs) != 1 || len(reservation.Receipts) != 0 { t.Fatal("invalid public initial reservation", err) } diff --git a/services/core/internal/store/session_artifacts_public_test.go b/services/core/internal/store/session_artifacts_public_test.go index 8447e86de..d5e6cc722 100644 --- a/services/core/internal/store/session_artifacts_public_test.go +++ b/services/core/internal/store/session_artifacts_public_test.go @@ -37,7 +37,7 @@ func hostedArtifactSession(t *testing.T, s *store.Store, tenant, key string) (se // through artifacts and settled as completed, and returns the Turn ID. func completeArtifactTurn(t *testing.T, s *store.Store, artifacts *sessions.Service, tenant, session, environment, key string, outputs map[string]string) string { t.Helper() - receipt, err := s.SubmitMessage(t.Context(), tenant, session, key, json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"publish outputs"}]}]}`)) + receipt, err := store.SendMessage(t.Context(), s, tenant, session, key, json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"publish outputs"}]}]}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_artifacts_test.go b/services/core/internal/store/session_artifacts_test.go index e35b125f7..ba5957f9f 100644 --- a/services/core/internal/store/session_artifacts_test.go +++ b/services/core/internal/store/session_artifacts_test.go @@ -235,7 +235,7 @@ func TestSessionArtifactTransferDoesNotBlockDeletionOrCancellation(t *testing.T) t.Fatalf("transfer blocked deletion: %v", err) } } else { - if _, err := s.RequestCancel(ctx, tenant, session, "cancel-capture"); err != nil { + if _, err := requestCancel(ctx, s, tenant, session, "cancel-capture"); err != nil { t.Fatalf("transfer blocked cancellation: %v", err) } want = sessions.ErrTurnConflict diff --git a/services/core/internal/store/session_creation_stream_test.go b/services/core/internal/store/session_creation_stream_test.go index fd0e70e19..c7aee3ae1 100644 --- a/services/core/internal/store/session_creation_stream_test.go +++ b/services/core/internal/store/session_creation_stream_test.go @@ -121,7 +121,7 @@ func TestCreationStreamStartsBeforeOwnInputsAndRetriesAtUpsertCursor(t *testing. if err != nil || late.Created || late.Cursor != all[len(all)-1].Sequence { t.Fatal(late, err) } - next, err := s.SubmitInputs(ctx, tenant, id, "next", []sessions.Input{messageInput("later")}) + next, err := submitInputs(ctx, s, tenant, id, "next", []sessions.Input{messageInput("later")}) if err != nil { t.Fatal(err) } @@ -156,7 +156,7 @@ func TestCreationStreamIdleAndNonstreamRetry(t *testing.T) { if err != nil || !first.Created || first.Cursor != 0 || first.Session.LastTurn != nil { t.Fatal(first, err) } - if _, err := s.SubmitInputs(ctx, tenant, first.Session.ID, "message", []sessions.Input{messageInput("later")}); err != nil { + if _, err := submitInputs(ctx, s, tenant, first.Session.ID, "message", []sessions.Input{messageInput("later")}); err != nil { t.Fatal(err) } retry, err := s.CreateSessionStream(ctx, tenant, input) diff --git a/services/core/internal/store/session_deletion_execution_test.go b/services/core/internal/store/session_deletion_execution_test.go index fe40df03a..92798a4b2 100644 --- a/services/core/internal/store/session_deletion_execution_test.go +++ b/services/core/internal/store/session_deletion_execution_test.go @@ -120,7 +120,7 @@ func TestWaitingSessionCancelsThenDeletesThroughWorker(t *testing.T) { if again := functionState(t, h, 1); again.LastTurn == nil || again.LastTurn.Status != sessions.TurnWaiting || !again.LastTurn.CancelRequestedAt.IsZero() { t.Fatal("rejected deletion changed required actions", again) } - if _, err := h.s.RequestCancel(ctx, h.tenant, h.session.ID, "cancel-before-delete"); err != nil { + if _, err := store.RequestCancel(ctx, h.s, h.tenant, h.session.ID, "cancel-before-delete"); err != nil { t.Fatal(err) } // The explicit cancellation, not the rejected deletion, reaches the daemon. diff --git a/services/core/internal/store/session_deletion_lifecycle_public_test.go b/services/core/internal/store/session_deletion_lifecycle_public_test.go index 04fd25a7f..8852cfda3 100644 --- a/services/core/internal/store/session_deletion_lifecycle_public_test.go +++ b/services/core/internal/store/session_deletion_lifecycle_public_test.go @@ -40,7 +40,6 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { defer server.Close() client := pathIDClient{t: t, server: server} leased := executionOwner(t, db, s) - writer := leased.Store create := func(environment string, initial bool) sessions.Session { t.Helper() @@ -61,7 +60,7 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { turn := func(to ...string) string { t.Helper() session := create(none, false) - receipt, err := s.SubmitMessage(ctx, tenant, session.ID, "input", json.RawMessage(`{"text":"work"}`)) + receipt, err := store.SendMessage(ctx, s, tenant, session.ID, "input", json.RawMessage(`{"text":"work"}`)) if err != nil { t.Fatal(err) } @@ -69,7 +68,7 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { for _, status := range to { switch status { case "cancel": - _, err = s.RequestCancel(ctx, tenant, session.ID, "cancel") + _, err = store.RequestCancel(ctx, s, tenant, session.ID, "cancel") case "function": err = leased.Sessions.RecordFunctionCall(ctx, tenant, session.ID, receipt.TurnID, sessions.FunctionCall{CallID: "pending", ExecutorCallID: "native-pending", Name: "lookup", Arguments: json.RawMessage(`{}`)}) case sessions.TurnCompleted, sessions.TurnFailed: @@ -86,7 +85,7 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { } reserve := func(session sessions.Session) sessions.EnvironmentInputReservation { t.Helper() - reservation, err := s.ReserveEnvironmentInput(ctx, tenant, session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}) + reservation, err := store.SessionService(t, s).ReserveEnvironmentInput(ctx, tenant, session.ID, "later", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"later"}`)}}) if err != nil || reservation.State != sessions.EnvironmentInputPending { t.Fatal(reservation, err) } @@ -139,12 +138,12 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { if _, err := db.pool.Exec(ctx, "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1", expired.ID); err != nil { t.Fatal(err) } - if count, err := writer.ExpireEnvironmentInputs(ctx); err != nil || count != 1 { + if count, err := leased.Sessions.ExpireEnvironmentInputs(ctx); err != nil || count != 1 { t.Fatal("initial input did not expire", count, err) } settled["self_hosted_input_expired"] = expired.ID withdrawn := create(selfHosted, false) - if _, err := s.CancelEnvironmentInput(ctx, tenant, withdrawn.ID, reserve(withdrawn).ID); err != nil { + if _, err := store.CancelEnvironmentInput(ctx, s, tenant, withdrawn.ID, reserve(withdrawn).ID); err != nil { t.Fatal(err) } settled["later_input_cancelled"] = withdrawn.ID diff --git a/services/core/internal/store/session_deletion_test.go b/services/core/internal/store/session_deletion_test.go index 04b1b5623..425307967 100644 --- a/services/core/internal/store/session_deletion_test.go +++ b/services/core/internal/store/session_deletion_test.go @@ -62,7 +62,7 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { if err != nil { t.Fatal(err) } - receipt, err := s.SubmitMessage(ctx, tenant, session.ID, "input", json.RawMessage(`{"text":"retained"}`)) + receipt, err := sendMessage(ctx, s, tenant, session.ID, "input", json.RawMessage(`{"text":"retained"}`)) if err != nil { t.Fatal(err) } @@ -101,7 +101,7 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { } // Callers cancel first. A queued Turn cancels at once; a running // Turn stays active until execution settles its cancellation. - if _, err := s.RequestCancel(ctx, tenant, session.ID, "cancel"); err != nil { + if _, err := requestCancel(ctx, s, tenant, session.ID, "cancel"); err != nil { t.Fatal(err) } if status == sessions.TurnInProgress { @@ -139,10 +139,10 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { if _, err := fresh.CreateSessionStream(ctx, tenant, input); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal(err) } - if _, err := fresh.SubmitMessage(ctx, tenant, session.ID, "input", json.RawMessage(`{"text":"retained"}`)); !errors.Is(err, sessions.ErrNotFound) { + if _, err := sendMessage(ctx, fresh, tenant, session.ID, "input", json.RawMessage(`{"text":"retained"}`)); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - if _, err := fresh.RequestCancel(ctx, tenant, session.ID, "late-cancel"); !errors.Is(err, sessions.ErrNotFound) { + if _, err := requestCancel(ctx, fresh, tenant, session.ID, "late-cancel"); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } if _, err := sessionAdapter(fresh).ListItems(ctx, tenant, session.ID, "", 20, true); !errors.Is(err, sessions.ErrNotFound) { @@ -162,7 +162,7 @@ func TestSessionDeletionWaitsForSettledTurnAndRejectsAdmission(t *testing.T) { if _, err := transitionTurn(ctx, fresh, tenant, session.ID, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal(err) } - inputs, err := fresh.ListTurnInputs(ctx, tenant, session.ID, receipt.TurnID, 0, 20) + inputs, err := sessionAdapter(fresh).ListTurnInputs(ctx, tenant, session.ID, receipt.TurnID, 0, 20) if err != nil || len(inputs) == 0 || inputs[0].Sequence != receipt.Sequence { t.Fatal(inputs, err) } @@ -184,7 +184,7 @@ func TestSessionDeletionSerializesAdmissionBeforeRetryLookup(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := s.RequestCancel(ctx, tenant, session.ID, "existing"); err != nil { + if _, err := requestCancel(ctx, s, tenant, session.ID, "existing"); err != nil { t.Fatal(err) } tx, err := pool.Begin(ctx) @@ -196,7 +196,7 @@ func TestSessionDeletionSerializesAdmissionBeforeRetryLookup(t *testing.T) { t.Fatal(err) } done := make(chan error, 1) - go func() { _, err := s.RequestCancel(ctx, tenant, session.ID, "existing"); done <- err }() + go func() { _, err := requestCancel(ctx, s, tenant, session.ID, "existing"); done <- err }() if err := tx.Commit(ctx); err != nil { t.Fatal(err) } @@ -262,14 +262,14 @@ func TestSessionDeletionRacesAdmissionUnderSessionLock(t *testing.T) { tenant, session := newTurnSession(t, s) return tenant, session.ID }, func(ctx context.Context, s *Store, tenant, session string) error { - _, err := s.SubmitMessage(ctx, tenant, session, "racing", messagePayload) + _, err := sendMessage(ctx, s, tenant, session, "racing", messagePayload) return err }}, {"environment_input", func(t *testing.T, s *Store) (string, string) { tenant, session := environmentInputSession(t, s) return tenant, session.ID }, func(ctx context.Context, s *Store, tenant, session string) error { - _, err := s.ReserveEnvironmentInput(ctx, tenant, session, "racing", []sessions.Input{messageInput("racing")}) + _, err := sessionService(t, s).ReserveEnvironmentInput(ctx, tenant, session, "racing", []sessions.Input{messageInput("racing")}) return err }}, } @@ -423,7 +423,7 @@ func TestSessionDeletionKeepsProvisioningInputPlacementUntilSettled(t *testing.T if _, err := s.pool.Exec(ctx, "UPDATE environment_input_reservations SET deadline=clock_timestamp()-interval '1 second' WHERE session_id=$1", session.ID); err != nil { t.Fatal(err) } - if count, err := w.ExpireEnvironmentInputs(ctx); err != nil || count != 1 { + if count, err := sessionExecution(t, w.lease).ExpireEnvironmentInputs(ctx); err != nil || count != 1 { t.Fatal("initial input did not expire", count, err) } if err := sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: session.ID}); err != nil { diff --git a/services/core/internal/store/session_environment_snapshot_test.go b/services/core/internal/store/session_environment_snapshot_test.go index 56e9ce389..eda9b4990 100644 --- a/services/core/internal/store/session_environment_snapshot_test.go +++ b/services/core/internal/store/session_environment_snapshot_test.go @@ -25,7 +25,7 @@ func TestSelfHostedCreationSnapshotRetainsEnvironmentAndCursor(t *testing.T) { if environment.ID == "" || environment.SessionID != created.Session.ID || environment.TenantID != tenant || environment.Status != "pending" { t.Fatal("incorrect creation association", environment) } - pending, err := s.ReserveEnvironmentInput(t.Context(), tenant, created.Session.ID, "later", []sessions.Input{messageInput("later")}) + pending, err := sessionService(t, s).ReserveEnvironmentInput(t.Context(), tenant, created.Session.ID, "later", []sessions.Input{messageInput("later")}) if err != nil || pending.State != sessions.EnvironmentInputPending { t.Fatal(pending, err) } diff --git a/services/core/internal/store/session_events_test.go b/services/core/internal/store/session_events_test.go index 3b7748f58..1c761eb11 100644 --- a/services/core/internal/store/session_events_test.go +++ b/services/core/internal/store/session_events_test.go @@ -57,7 +57,7 @@ func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(ctx, tenant, session.ID, "start", json.RawMessage(`{"text":"question"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "start", json.RawMessage(`{"text":"question"}`)) if err != nil { t.Fatal(err) } @@ -65,7 +65,7 @@ func TestSessionEventsCommitSnapshotsRetriesAndIsolation(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err = s.SubmitMessage(ctx, tenant, session.ID, "start", json.RawMessage(`{"text":"question"}`)); err != nil { + if _, err = store.SendMessage(ctx, s, tenant, session.ID, "start", json.RawMessage(`{"text":"question"}`)); err != nil { t.Fatal(err) } after, _ := store.SessionAdapter(s).SessionEventCursor(ctx, tenant, session.ID) @@ -170,7 +170,7 @@ func TestSessionEventsRetentionAndQueuedCancellation(t *testing.T) { inputs[i] = sessions.Input{Kind: "message", Payload: json.RawMessage(`{"text":"input"}`)} } for range 5 { - if _, err = s.SubmitInputs(ctx, tenant, session.ID, uuid.NewString(), inputs); err != nil { + if _, err = store.SubmitInputs(ctx, s, tenant, session.ID, uuid.NewString(), inputs); err != nil { t.Fatal(err) } } @@ -185,7 +185,7 @@ func TestSessionEventsRetentionAndQueuedCancellation(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err = s.RequestCancel(ctx, tenant, session.ID, "cancel"); err != nil { + if _, err = store.RequestCancel(ctx, s, tenant, session.ID, "cancel"); err != nil { t.Fatal(err) } changes, err := store.SessionAdapter(s).ListSessionEvents(ctx, tenant, session.ID, cursor) diff --git a/services/core/internal/store/session_initial_input.go b/services/core/internal/store/session_initial_input.go index 15c9642ed..77e89638b 100644 --- a/services/core/internal/store/session_initial_input.go +++ b/services/core/internal/store/session_initial_input.go @@ -21,15 +21,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/writeaudit" ) -func validateInitialInputs(inputs []sessions.Input) ([]sessions.Input, json.RawMessage, error) { - for _, input := range inputs { - if input.Kind != "message" { - return nil, nil, sessions.ErrInvalidInput - } - } - return validateInputs(inputs) -} - // The Session upsert locks retries. Only the new row reserves or admits work, so a // retry after completion or later Turns cannot submit the original input again. func (s *Store) createSessionResources(ctx context.Context, tenant string, params sqlc.CreateSessionParams, inputs []sessions.Input, encodedInput json.RawMessage, files []environmentconfig.InitialFile, setup environmentconfig.Setup, provider *v1.ModelProviderInput, executionConfiguration *v1.SessionExecutionConfiguration, providerSource string, deploymentRevision uuid.UUID) (sqlc.Session, *sessions.Environment, error) { @@ -170,10 +161,8 @@ func (s *Store) createSessionResources(ctx context.Context, tenant string, param return err } } else { - for position, input := range inputs { - if _, err := admitInput(ctx, q, tenant, row.ID, key, int32(position), input); err != nil { - return err - } + if _, err := sessions.AdmitInputs(ctx, sessionpg.BindSession(q, row.TenantID, row.ID), key, inputs); err != nil { + return err } } if err := sessionpg.PruneChanges(ctx, q, row.ID); err != nil { diff --git a/services/core/internal/store/session_initial_input_test.go b/services/core/internal/store/session_initial_input_test.go index 040701873..d8d568786 100644 --- a/services/core/internal/store/session_initial_input_test.go +++ b/services/core/internal/store/session_initial_input_test.go @@ -50,7 +50,7 @@ func TestInitialInputCreationRetriesAcrossConnectionsAndLaterTurns(t *testing.T) if first.LastTurn == nil { t.Fatal("missing initial Turn") } - inputs, err := s.ListTurnInputs(ctx, tenant, first.ID, first.LastTurn.ID, 0, 100) + inputs, err := sessionAdapter(s).ListTurnInputs(ctx, tenant, first.ID, first.LastTurn.ID, 0, 100) if err != nil || len(inputs) != 2 { t.Fatal(inputs, err) } @@ -67,11 +67,11 @@ func TestInitialInputCreationRetriesAcrossConnectionsAndLaterTurns(t *testing.T) transition(t, s, tenant, first.ID, first.LastTurn.ID, sessions.TurnQueued, sessions.TurnInProgress) transition(t, s, tenant, first.ID, first.LastTurn.ID, sessions.TurnInProgress, sessions.TurnCompleted) // The same caller key at the events endpoint is an independent request. - next, err := s.SubmitInputs(ctx, tenant, first.ID, input.IdempotencyKey, []sessions.Input{messageInput("later")}) + next, err := submitInputs(ctx, s, tenant, first.ID, input.IdempotencyKey, []sessions.Input{messageInput("later")}) if err != nil || len(next) != 1 || next[0].TurnID == first.LastTurn.ID { t.Fatal(next, err) } - if _, err := s.RequestCancel(ctx, tenant, first.ID, "cancel"); err != nil { + if _, err := requestCancel(ctx, s, tenant, first.ID, "cancel"); err != nil { t.Fatal(err) } if _, err := sessionService(t, s).UpdateSessionMetadata(ctx, sessions.UpdateSessionMetadataCommand{TenantID: tenant, SessionID: first.ID, Metadata: map[string]string{"updated": "yes"}}); err != nil { diff --git a/services/core/internal/store/session_metadata_test.go b/services/core/internal/store/session_metadata_test.go index a506b7b2a..d1be257f1 100644 --- a/services/core/internal/store/session_metadata_test.go +++ b/services/core/internal/store/session_metadata_test.go @@ -114,7 +114,7 @@ func TestSessionMetadataPreservesTerminalActivity(t *testing.T) { t.Fatal(err) } for _, status := range []string{sessions.TurnCompleted, sessions.TurnFailed, sessions.TurnCancelled} { - receipt, err := s.SubmitMessage(ctx, tenant, session.ID, uuid.NewString(), []byte(`{"text":"metadata fixture"}`)) + receipt, err := sendMessage(ctx, s, tenant, session.ID, uuid.NewString(), []byte(`{"text":"metadata fixture"}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/session_write_audit_test.go b/services/core/internal/store/session_write_audit_test.go index fef29312b..61249e770 100644 --- a/services/core/internal/store/session_write_audit_test.go +++ b/services/core/internal/store/session_write_audit_test.go @@ -143,11 +143,11 @@ func TestSessionWriteAuditRollback(t *testing.T) { case "delete": err = sessionService(t, s).DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: tenant, SessionID: created.ID}) case "events": - _, err = s.SubmitInputs(ctx, tenant, created.ID, "events", []sessions.Input{messageInput("private")}) + _, err = submitInputs(ctx, s, tenant, created.ID, "events", []sessions.Input{messageInput("private")}) case "noop": err = sessionService(t, s).AuditSessionOperation(ctx, sessions.AuditSessionOperationCommand{TenantID: tenant, SessionID: created.ID, Action: "send_events"}) case "reserve": - _, err = s.ReserveEnvironmentInput(ctx, tenant, created.ID, "reserve", []sessions.Input{messageInput("private")}) + _, err = sessionService(t, s).ReserveEnvironmentInput(ctx, tenant, created.ID, "reserve", []sessions.Input{messageInput("private")}) } if err == nil { t.Fatal("mutation bypassed audit failure") @@ -177,10 +177,10 @@ func TestEventsWriteAuditAdmissionAndReplay(t *testing.T) { } submit := func(ctx context.Context) error { if prepared { - _, err := s.ReserveEnvironmentInput(ctx, tenant, created.ID, "batch", []sessions.Input{messageInput("private")}) + _, err := sessionService(t, s).ReserveEnvironmentInput(ctx, tenant, created.ID, "batch", []sessions.Input{messageInput("private")}) return err } - _, err := s.SubmitInputs(ctx, tenant, created.ID, "batch", []sessions.Input{messageInput("private")}) + _, err := submitInputs(ctx, s, tenant, created.ID, "batch", []sessions.Input{messageInput("private")}) return err } first := sessionAuditContext(t, tenant, "a") diff --git a/services/core/internal/store/sessions.go b/services/core/internal/store/sessions.go index 95c89ba13..78dcf1f94 100644 --- a/services/core/internal/store/sessions.go +++ b/services/core/internal/store/sessions.go @@ -89,7 +89,7 @@ func (s *Store) createSession(ctx context.Context, tenantID string, input sessio var batch []sessions.Input var encodedInput json.RawMessage if len(input.InitialInputs) > 0 { - batch, encodedInput, err = validateInitialInputs(input.InitialInputs) + batch, encodedInput, err = sessions.ValidateMessageInputs(input.InitialInputs) if err != nil { return sessions.Creation{}, err } diff --git a/services/core/internal/store/steering_receipts_test.go b/services/core/internal/store/steering_receipts_test.go index beff34513..69a7ab13e 100644 --- a/services/core/internal/store/steering_receipts_test.go +++ b/services/core/internal/store/steering_receipts_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" ) func TestExecutionDurableInputReceiptLifetime(t *testing.T) { @@ -42,7 +43,7 @@ func TestExecutionDurableInputReceiptLifetime(t *testing.T) { case "retry-after-write": h.write(first.TurnID, proto.TypePromptSteerAck, proto.PromptSteerAckPayload{InputID: input.InputID, ErrorCode: "not_ready"}) case "cancel-unknown-first": - if _, err := h.s.RequestCancel(ctx, h.tenant, h.session.ID, "cancel"); err != nil { + if _, err := store.RequestCancel(ctx, h.s, h.tenant, h.session.ID, "cancel"); err != nil { t.Fatal(err) } var request proto.PromptCancelPayload diff --git a/services/core/internal/store/stream_authority_http_test.go b/services/core/internal/store/stream_authority_http_test.go index a03a3fc5c..1d68d49eb 100644 --- a/services/core/internal/store/stream_authority_http_test.go +++ b/services/core/internal/store/stream_authority_http_test.go @@ -101,7 +101,7 @@ func TestLiveStreamClosesAfterKeyRevocationOrProjectArchive(t *testing.T) { t.Fatal("revocation affected peer", valid.StatusCode) } // New events remain available to valid callers after the reader has closed. - if _, err := s.ReserveEnvironmentInput(t.Context(), project.TenantID, session.ID, "after-revocation", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"new event"}`)}}); err != nil { + if _, err := store.SessionService(t, s).ReserveEnvironmentInput(t.Context(), project.TenantID, session.ID, "after-revocation", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"new event"}`)}}); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/token_usage_integration_test.go b/services/core/internal/store/token_usage_integration_test.go index a042c21c2..dfddb68fd 100644 --- a/services/core/internal/store/token_usage_integration_test.go +++ b/services/core/internal/store/token_usage_integration_test.go @@ -34,7 +34,7 @@ func TestTokenUsageDurableSnapshotsAndSessionTotals(t *testing.T) { } } for n, status := range []string{sessions.TurnFailed, sessions.TurnCancelled} { - admission, err := s.SubmitMessage(ctx, tenant, session.ID, fmt.Sprint(n), json.RawMessage(`{"text":"measure"}`)) + admission, err := store.SendMessage(ctx, s, tenant, session.ID, fmt.Sprint(n), json.RawMessage(`{"text":"measure"}`)) if err != nil { t.Fatal(err) } @@ -114,7 +114,7 @@ func TestCancellationReceiptUsageSurvivesRecovery(t *testing.T) { if err != nil { t.Fatal(err) } - admission, err := s.SubmitMessage(ctx, tenant, session.ID, "start", json.RawMessage(`{"text":"measure"}`)) + admission, err := store.SendMessage(ctx, s, tenant, session.ID, "start", json.RawMessage(`{"text":"measure"}`)) if err != nil { t.Fatal(err) } @@ -188,7 +188,7 @@ func TestSessionUsageRequiresEveryRootTurnEndedAndMeasured(t *testing.T) { } submit := func(key string) sessions.InputReceipt { t.Helper() - admission, err := s.SubmitMessage(ctx, tenant, session.ID, key, json.RawMessage(`{"text":"measure"}`)) + admission, err := store.SendMessage(ctx, s, tenant, session.ID, key, json.RawMessage(`{"text":"measure"}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/tool_policy_native_test.go b/services/core/internal/store/tool_policy_native_test.go index bfa83bfc0..380a0c4f7 100644 --- a/services/core/internal/store/tool_policy_native_test.go +++ b/services/core/internal/store/tool_policy_native_test.go @@ -95,7 +95,7 @@ func TestNativeToolPolicyPublicExecution(t *testing.T) { if err != nil { t.Fatal(err) } - inputs, err := h.s.ListTurnInputs(ctx, h.tenant, item.ID, item.FirstTurn, 0, 100) + inputs, err := store.SessionAdapter(h.s).ListTurnInputs(ctx, h.tenant, item.ID, item.FirstTurn, 0, 100) var outcome execution.Result if err != nil || len(inputs) != 1 || json.Unmarshal(turn.Outcome, &outcome) != nil || outcome.AppliedThrough != inputs[0].Sequence { t.Fatal("native text input receipt missing", err) diff --git a/services/core/internal/store/turn_events_test.go b/services/core/internal/store/turn_events_test.go index bd4c1c9ec..0e8b8ba0c 100644 --- a/services/core/internal/store/turn_events_test.go +++ b/services/core/internal/store/turn_events_test.go @@ -21,7 +21,7 @@ func TestTurnEventBatchesAreOrderedIsolatedAndDurable(t *testing.T) { if err != nil { t.Fatal(err) } - input, err := s.SubmitMessage(ctx, tenant, session.ID, "start", json.RawMessage(`{"text":"test"}`)) + input, err := store.SendMessage(ctx, s, tenant, session.ID, "start", json.RawMessage(`{"text":"test"}`)) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/turn_inputs.go b/services/core/internal/store/turn_inputs.go deleted file mode 100644 index 3e896dafd..000000000 --- a/services/core/internal/store/turn_inputs.go +++ /dev/null @@ -1,238 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "slices" - - "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/jsonobject" - "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" -) - -// SubmitMessage and RequestCancel use the same request-level admission as batches. -func (s *Store) SubmitMessage(ctx context.Context, tenantID, sessionID, key string, payload json.RawMessage) (sessions.InputReceipt, error) { - return s.submitOne(ctx, tenantID, sessionID, key, sessions.Input{Kind: "message", Payload: payload}) -} - -func (s *Store) RequestCancel(ctx context.Context, tenantID, sessionID, key string) (sessions.InputReceipt, error) { - return s.submitOne(ctx, tenantID, sessionID, key, sessions.Input{Kind: "cancel", Payload: json.RawMessage(`{}`)}) -} - -func (s *Store) submitOne(ctx context.Context, tenantID, sessionID, key string, input sessions.Input) (sessions.InputReceipt, error) { - receipts, err := s.SubmitInputs(ctx, tenantID, sessionID, key, []sessions.Input{input}) - if err != nil { - return sessions.InputReceipt{}, err - } - return receipts[0], nil -} - -// SubmitInputs commits a request in order under one Session lock. The entire -// batch is the retry identity; replay never re-evaluates a cancellation target. -// Internal receipts are not the response body of the public events endpoint. -func (s *Store) SubmitInputs(ctx context.Context, tenantID, sessionID, key string, inputs []sessions.Input) ([]sessions.InputReceipt, error) { - if err := sessions.ValidateInputKey(key); err != nil { - return nil, err - } - batch, encoded, err := validateInputs(inputs) - if err != nil { - return nil, err - } - tenant, err := parseID(tenantID) - if err != nil { - return nil, err - } - receipts := make([]sessions.InputReceipt, 0, len(batch)) - err = s.withPublicSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - previous, err := inputBatchReceipts(ctx, q, session, key, encoded) - if err != nil { - return err - } - if len(previous) > 0 { - receipts = previous - return auditpg.RecordWriteAudit(ctx, q, tenantID, "send_events", "session", uuid.UUID(session.Bytes).String(), "") - } - if slices.ContainsFunc(batch, func(input sessions.Input) bool { return input.Kind == "message" }) { - if err := sessions.CheckFileWriteGate(ctx, sessionpg.BindSession(q, tenant, session)); err != nil { - return err - } - } - if err := checkEnvironmentInputGate(ctx, q, session, key, encoded); err != nil { - return err - } - for position, input := range batch { - receipt, err := admitInput(ctx, q, tenantID, session, key, int32(position), input) - if err != nil { - return err - } - receipts = append(receipts, receipt) - } - return auditpg.RecordWriteAudit(ctx, q, tenantID, "send_events", "session", uuid.UUID(session.Bytes).String(), "") - }) - if err != nil { - return nil, fmt.Errorf("submit turn inputs: %w", err) - } - return receipts, nil -} - -func inputBatchReceipts(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, key string, batch json.RawMessage) ([]sessions.InputReceipt, error) { - rows, err := q.FindInputBatch(ctx, sqlc.FindInputBatchParams{SessionID: session, IdempotencyKey: key, Batch: batch}) - if err != nil { - return nil, err - } - receipts := make([]sessions.InputReceipt, 0, len(rows)) - for _, row := range rows { - if !row.Matches { - return nil, sessions.ErrIdempotencyConflict - } - receipts = append(receipts, inputReceipt(row.Sequence, row.TurnID, true)) - } - return receipts, nil -} - -func validateInputs(inputs []sessions.Input) ([]sessions.Input, json.RawMessage, error) { - if len(inputs) == 0 || len(inputs) > 64 { - return nil, nil, fmt.Errorf("%w: input batch must contain 1..64 events", sessions.ErrInvalidInput) - } - batch := make([]sessions.Input, len(inputs)) - size := 0 - for i, input := range inputs { - size += len(input.Payload) - if size > 512*1024 || len(input.Payload) == 0 || (input.Kind != "message" && input.Kind != "cancel" && input.Kind != "tool_result") { - return nil, nil, fmt.Errorf("%w: input payloads must be nonempty and total at most 512 KiB", sessions.ErrInvalidInput) - } - payload, err := jsonobject.Normalize(input.Payload) - if err != nil { - return nil, nil, fmt.Errorf("%w: %w", sessions.ErrInvalidInput, err) - } - if input.Kind == "cancel" && string(payload) != "{}" { - return nil, nil, fmt.Errorf("%w: cancel payload must be empty", sessions.ErrInvalidInput) - } - if input.Kind == "tool_result" { - if _, err := sessions.ParseFunctionResultInput(payload); err != nil { - return nil, nil, err - } - } - batch[i] = sessions.Input{Kind: input.Kind, Payload: payload} - } - encoded, err := json.Marshal(batch) - return batch, encoded, err -} - -func admitInput(ctx context.Context, q *sqlc.Queries, tenantID string, session pgtype.UUID, key string, position int32, input sessions.Input) (sessions.InputReceipt, error) { - if input.Kind == "tool_result" { - result, err := sessions.ParseFunctionResultInput(input.Payload) - if err != nil { - return sessions.InputReceipt{}, err - } - tenant, err := parseID(tenantID) - if err != nil { - return sessions.InputReceipt{}, err - } - turn, err := sessions.AdmitFunctionResult(ctx, sessionpg.BindSession(q, tenant, session), result) - if err != nil { - return sessions.InputReceipt{}, err - } - id, err := parseID(turn.ID) - if err != nil { - return sessions.InputReceipt{}, err - } - sequence, err := q.CreateTurnInput(ctx, sqlc.CreateTurnInputParams{ - SessionID: session, TurnID: id, IdempotencyKey: key, Kind: input.Kind, Payload: input.Payload, BatchPosition: position, - }) - if err != nil { - return sessions.InputReceipt{}, err - } - return inputReceipt(sequence, id, false), nil - } - created := false - turn, err := q.GetActiveTurn(ctx, session) - if errors.Is(err, pgx.ErrNoRows) { - if input.Kind == "message" { - turn, err = q.CreateTurn(ctx, sqlc.CreateTurnParams{ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, SessionID: session}) - if err == nil { - created = true - err = sessionpg.AppendChanges(ctx, q, session, sessions.TurnChanges(sessionpg.TurnFromRow(turn), true)...) - } - } else { - err = nil // Retain even an idle cancellation's retry identity. - } - } - if err != nil { - return sessions.InputReceipt{}, err - } - sequence, err := q.CreateTurnInput(ctx, sqlc.CreateTurnInputParams{ - SessionID: session, TurnID: turn.ID, IdempotencyKey: key, Kind: input.Kind, Payload: input.Payload, BatchPosition: position, - }) - if err != nil { - return sessions.InputReceipt{}, err - } - tenant, err := parseID(tenantID) - if err != nil { - return sessions.InputReceipt{}, err - } - bound := sessionpg.BindSession(q, tenant, session) - if input.Kind == "cancel" && turn.ID.Valid { - if err := sessions.CancelTurn(ctx, bound, sessionpg.TurnFromRow(turn)); err != nil { - return sessions.InputReceipt{}, err - } - } - if err := sessions.ProjectInput(ctx, bound, sequence); err != nil { - return sessions.InputReceipt{}, err - } - if created { - // A new Turn publishes turn.created, then its user input Items, then the - // Session activity, within this transaction. - usage, err := sessionpg.LoadUsage(ctx, q, session) - if err != nil { - return sessions.InputReceipt{}, err - } - if err := sessionpg.AppendChanges(ctx, q, session, sessions.ActivityChange(sessionpg.TurnFromRow(turn), usage, nil)); err != nil { - return sessions.InputReceipt{}, err - } - } - return inputReceipt(sequence, turn.ID, false), nil -} - -// ListTurnInputs is an internal ordered recovery query, not the public SSE stream. -func (s *Store) ListTurnInputs(ctx context.Context, tenantID, sessionID, turnID string, after int64, limit int) ([]sessions.TurnInput, error) { - params, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) - if err != nil { - return nil, err - } - if after < 0 || limit < 1 || limit > 100 { - return nil, fmt.Errorf("%w: nonnegative cursor and page size 1..100 required", sessions.ErrInvalidInput) - } - if _, err := s.queries.GetTurn(ctx, params); errors.Is(err, pgx.ErrNoRows) { - return nil, sessions.ErrNotFound - } else if err != nil { - return nil, fmt.Errorf("get turn: %w", err) - } - rows, err := s.queries.ListTurnInputs(ctx, sqlc.ListTurnInputsParams{ - TenantID: params.TenantID, SessionID: params.SessionID, TurnID: params.ID, Sequence: after, Limit: int32(limit), - }) - if err != nil { - return nil, fmt.Errorf("list turn inputs: %w", err) - } - inputs := make([]sessions.TurnInput, 0, len(rows)) - for _, row := range rows { - inputs = append(inputs, sessions.TurnInput{Sequence: row.Sequence, Kind: row.Kind, Payload: row.Payload, CreatedAt: row.CreatedAt.Time}) - } - return inputs, nil -} - -func inputReceipt(sequence int64, turn pgtype.UUID, replayed bool) sessions.InputReceipt { - receipt := sessions.InputReceipt{Sequence: sequence, Replayed: replayed} - if turn.Valid { - receipt.TurnID = uuid.UUID(turn.Bytes).String() - } - return receipt -} diff --git a/services/core/internal/store/turn_inputs_test.go b/services/core/internal/store/turn_inputs_test.go deleted file mode 100644 index 3ffc28e26..000000000 --- a/services/core/internal/store/turn_inputs_test.go +++ /dev/null @@ -1,219 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "reflect" - "sync" - "testing" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" -) - -var messagePayload = json.RawMessage(`{"input":[{"role":"user","content":[{"type":"input_text","text":"hello"}]}]}`) - -func newTurnSession(t *testing.T, s *Store) (string, sessions.Session) { - t.Helper() - tenant := uuid.NewString() - session, err := s.CreateSession(context.Background(), tenant, sessions.CreateSession{Creator: FixtureCreator(), - Engine: "codex", IdempotencyKey: "session", Configuration: json.RawMessage(`{"agent":{"model":"test","instructions":"original"}}`), - }) - if err != nil { - t.Fatal(err) - } - return tenant, session -} - -func submitMessage(t *testing.T, s *Store, tenant, session, key string) sessions.InputReceipt { - t.Helper() - receipt, err := s.SubmitMessage(context.Background(), tenant, session, key, messagePayload) - if err != nil { - t.Fatal(err) - } - return receipt -} - -func transition(t *testing.T, s *Store, tenant, session, turn, from, to string) sessions.Turn { - t.Helper() - got, err := transitionTurn(context.Background(), s, tenant, session, turn, sessions.TurnTransition{ExpectedStatus: from, Status: to}) - if err != nil { - t.Fatal(err) - } - return got -} - -func TestConcurrentInputsUseOneTurnAndOneRetryReceipt(t *testing.T) { - s, _ := testStore(t) - otherStore, _ := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - const count = 8 - for _, repeatedKey := range []bool{true, false} { - t.Run(fmt.Sprintf("repeated-key-%v", repeatedKey), func(t *testing.T) { - var wg sync.WaitGroup - receipts := make(chan sessions.InputReceipt, count) - errs := make(chan error, count) - for i := range count { - wg.Add(1) - go func() { - defer wg.Done() - key := "same" - if !repeatedKey { - key = fmt.Sprintf("distinct-%d", i) - } - store := s - if i%2 == 0 { - store = otherStore - } - r, err := store.SubmitMessage(ctx, tenant, session.ID, key, messagePayload) - receipts <- r - errs <- err - }() - } - wg.Wait() - close(receipts) - close(errs) - for err := range errs { - if err != nil { - t.Fatal(err) - } - } - turns, sequences := map[string]bool{}, map[int64]bool{} - newReceipts := 0 - for r := range receipts { - turns[r.TurnID], sequences[r.Sequence] = true, true - if !r.Replayed { - newReceipts++ - } - } - want := count - if repeatedKey { - want = 1 - } - if len(turns) != 1 || turns[""] || len(sequences) != want || newReceipts != want { - t.Fatalf("turns=%v sequences=%v new=%d", turns, sequences, newReceipts) - } - }) - } - first := submitMessage(t, s, tenant, session.ID, "same") - inputs, err := s.ListTurnInputs(ctx, tenant, session.ID, first.TurnID, 0, 100) - if err != nil || len(inputs) != count+1 { - t.Fatalf("lost or duplicated inputs: %d, %v", len(inputs), err) - } -} - -func TestTurnInputRetriesAndRestart(t *testing.T) { - s, pool := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - first := submitMessage(t, s, tenant, session.ID, "first") - transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - steer := submitMessage(t, s, tenant, session.ID, "steer") - if steer.TurnID != first.TurnID || steer.Sequence <= first.Sequence { - t.Fatalf("active message did not steer: %+v", steer) - } - reordered := json.RawMessage(`{ "input": [{"content":[{"text":"hello","type":"input_text"}],"role":"user"}] }`) - retry, err := s.SubmitMessage(ctx, tenant, session.ID, "first", reordered) - if err != nil || !retry.Replayed || retry.Sequence != first.Sequence || retry.TurnID != first.TurnID { - t.Fatalf("equivalent retry = %+v, %v", retry, err) - } - if _, err := s.SubmitMessage(ctx, tenant, session.ID, "first", json.RawMessage(`{"text":"changed"}`)); !errors.Is(err, sessions.ErrIdempotencyConflict) { - t.Fatalf("changed payload accepted: %v", err) - } - if _, err := s.RequestCancel(ctx, tenant, session.ID, "first"); !errors.Is(err, sessions.ErrIdempotencyConflict) { - t.Fatalf("changed input kind accepted: %v", err) - } - completed := transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnInProgress, sessions.TurnCompleted) - next := submitMessage(t, s, tenant, session.ID, "next") - if next.TurnID == first.TurnID { - t.Fatal("idle message did not start a new Turn") - } - pool.Close() - recovered, _ := testStore(t) - retry = submitMessage(t, recovered, tenant, session.ID, "first") - if !retry.Replayed || retry.TurnID != first.TurnID || retry.Sequence != first.Sequence { - t.Fatalf("restart retry changed target: %+v", retry) - } - got, err := sessionAdapter(recovered).GetTurn(ctx, tenant, session.ID, first.TurnID) - if err != nil || !reflect.DeepEqual(got, completed) { - t.Fatalf("restart turn: %+v, %v", got, err) - } - var all []sessions.TurnInput - var cursor int64 - for { - page, err := recovered.ListTurnInputs(ctx, tenant, session.ID, first.TurnID, cursor, 1) - if err != nil { - t.Fatal(err) - } - if len(page) == 0 { - break - } - if page[0].Sequence <= cursor || page[0].CreatedAt.IsZero() { - t.Fatalf("invalid ordered input: %+v", page) - } - cursor = page[0].Sequence - all = append(all, page...) - } - if len(all) != 2 || all[0].Sequence != first.Sequence || all[1].Sequence != steer.Sequence { - t.Fatalf("recovered inputs = %+v", all) - } - snapshot, err := sessionAdapter(recovered).GetSession(ctx, tenant, session.ID) - if err != nil || snapshot.LastTurn == nil || snapshot.LastTurn.ID != next.TurnID || snapshot.LastTurn.Status != sessions.TurnQueued { - t.Fatal("latest Session activity did not survive restart", err) - } - snapshot.LastTurn = nil - if !reflect.DeepEqual(snapshot, session) { - t.Fatal("turn submission mutated the Session snapshot", err) - } -} - -func TestTurnOperationsAreTenantAndSessionScoped(t *testing.T) { - s, _ := testStore(t) - tenant, session := newTurnSession(t, s) - otherTenant, other := newTurnSession(t, s) - ctx := context.Background() - first := submitMessage(t, s, tenant, session.ID, "input") - for _, scope := range []struct{ tenant, session string }{{otherTenant, session.ID}, {tenant, other.ID}, {tenant, uuid.NewString()}} { - for name, call := range map[string]func() error{ - "submit": func() error { - _, err := s.SubmitMessage(ctx, scope.tenant, scope.session, "input", messagePayload) - return err - }, - "cancel": func() error { _, err := s.RequestCancel(ctx, scope.tenant, scope.session, "cancel"); return err }, - "read": func() error { - _, err := sessionAdapter(s).GetTurn(ctx, scope.tenant, scope.session, first.TurnID) - return err - }, - "inputs": func() error { - _, err := s.ListTurnInputs(ctx, scope.tenant, scope.session, first.TurnID, 0, 10) - return err - }, - "transition": func() error { - _, err := transitionTurn(ctx, s, scope.tenant, scope.session, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed}) - return err - }, - } { - if err := call(); !errors.Is(err, sessions.ErrNotFound) { - t.Fatalf("%s escaped scope: %v", name, err) - } - } - } - // Turn IDs cannot be used with another valid Session in the same tenant either. - second, err := s.CreateSession(ctx, tenant, sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "second"}) - if err != nil { - t.Fatal(err) - } - if _, err := sessionAdapter(s).GetTurn(ctx, tenant, second.ID, first.TurnID); !errors.Is(err, sessions.ErrNotFound) { - t.Fatalf("cross-session turn read: %v", err) - } - if _, err := transitionTurn(ctx, s, tenant, second.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnFailed}); !errors.Is(err, sessions.ErrNotFound) { - t.Fatalf("cross-session turn write: %v", err) - } - otherInput := submitMessage(t, s, otherTenant, other.ID, "input") - if otherInput.TurnID == first.TurnID { - t.Fatal("retry identity leaked across Sessions") - } -} diff --git a/services/core/internal/store/turns_test.go b/services/core/internal/store/turns_test.go deleted file mode 100644 index 4e85732e1..000000000 --- a/services/core/internal/store/turns_test.go +++ /dev/null @@ -1,114 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "testing" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" -) - -func TestCancellationStaysBoundToItsOriginalTurn(t *testing.T) { - s, _ := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - idle, err := s.RequestCancel(ctx, tenant, session.ID, "idle-cancel") - if err != nil || idle.TurnID != "" || idle.Sequence == 0 { - t.Fatalf("idle cancellation: %+v, %v", idle, err) - } - first := submitMessage(t, s, tenant, session.ID, "first") - started := transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - if started.StartedAt.IsZero() || !started.CompletedAt.IsZero() { - t.Fatalf("started timestamps: %+v", started) - } - cancel, err := s.RequestCancel(ctx, tenant, session.ID, "cancel") - if err != nil || cancel.TurnID != first.TurnID { - t.Fatalf("cancellation target: %+v, %v", cancel, err) - } - pending, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, first.TurnID) - if err != nil || pending.Status != sessions.TurnInProgress || pending.CancelRequestedAt.IsZero() || !pending.CompletedAt.IsZero() { - t.Fatalf("running cancellation falsely completed: %+v, %v", pending, err) - } - transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnInProgress, sessions.TurnCancelled) - next := submitMessage(t, s, tenant, session.ID, "next") - for key, original := range map[string]sessions.InputReceipt{"cancel": cancel, "idle-cancel": idle} { - retry, err := s.RequestCancel(ctx, tenant, session.ID, key) - if err != nil || !retry.Replayed || retry.TurnID != original.TurnID || retry.Sequence != original.Sequence { - t.Fatalf("cancellation retargeted: %+v, %v", retry, err) - } - } - queued, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, next.TurnID) - if err != nil || queued.Status != sessions.TurnQueued || !queued.CancelRequestedAt.IsZero() { - t.Fatalf("old cancellation affected later Turn: %+v, %v", queued, err) - } - if _, err := s.RequestCancel(ctx, tenant, session.ID, "cancel-queued"); err != nil { - t.Fatal(err) - } - stopped, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, next.TurnID) - if err != nil || stopped.Status != sessions.TurnCancelled || stopped.CompletedAt.IsZero() || !stopped.StartedAt.IsZero() { - t.Fatalf("queued work did not stop: %+v, %v", stopped, err) - } - if _, err := transitionTurn(ctx, s, tenant, session.ID, next.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatalf("cancelled queued work was started: %v", err) - } -} - -func TestWaitingTurnRetainsInputsAndStartTime(t *testing.T) { - s, _ := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - first := submitMessage(t, s, tenant, session.ID, "first") - started := transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnInProgress, sessions.TurnWaiting) - steer := submitMessage(t, s, tenant, session.ID, "steer") - if steer.TurnID != first.TurnID { - t.Fatal("waiting input started another Turn") - } - resumed := transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnWaiting, sessions.TurnInProgress) - if !started.StartedAt.Equal(resumed.StartedAt) { - t.Fatal("resume reset the start time") - } - transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnInProgress, sessions.TurnWaiting) - if _, err := s.RequestCancel(ctx, tenant, session.ID, "cancel"); err != nil { - t.Fatal(err) - } - if _, err := transitionTurn(ctx, s, tenant, session.ID, first.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnWaiting, Status: sessions.TurnInProgress}); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatalf("cancelling Turn resumed: %v", err) - } - transition(t, s, tenant, session.ID, first.TurnID, sessions.TurnWaiting, sessions.TurnCancelled) -} - -func TestTurnInputValidationHasNoSideEffects(t *testing.T) { - s, _ := testStore(t) - tenant, session := newTurnSession(t, s) - ctx := context.Background() - for _, raw := range []json.RawMessage{nil, json.RawMessage(`[]`), json.RawMessage(`null`), json.RawMessage(`{} {}`), json.RawMessage(`{"text":"` + string(make([]byte, 512*1024)) + `"}`)} { - if _, err := s.SubmitMessage(ctx, tenant, session.ID, "first", raw); !errors.Is(err, sessions.ErrInvalidInput) { - t.Fatalf("invalid input accepted: %v", err) - } - } - cancelled, stop := context.WithCancel(ctx) - stop() - if _, err := s.SubmitMessage(cancelled, tenant, session.ID, "first", messagePayload); !errors.Is(err, context.Canceled) { - t.Fatalf("cancelled context: %v", err) - } - first := submitMessage(t, s, tenant, session.ID, "first") - if first.Replayed { - t.Fatal("failed submission persisted a receipt") - } - for _, input := range []sessions.TurnTransition{ - {ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnCompleted}, - {ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnQueued}, - {ExpectedStatus: sessions.TurnCompleted, Status: sessions.TurnInProgress}, - {ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress, Outcome: json.RawMessage(`{"premature":true}`)}, - } { - if _, err := transitionTurn(ctx, s, tenant, session.ID, first.TurnID, input); !errors.Is(err, sessions.ErrInvalidInput) { - t.Fatalf("invalid transition accepted: %v", err) - } - } - got, err := sessionAdapter(s).GetTurn(ctx, tenant, session.ID, first.TurnID) - if err != nil || got.Status != sessions.TurnQueued || !got.StartedAt.IsZero() { - t.Fatalf("invalid transition changed Turn: %+v, %v", got, err) - } -} diff --git a/services/core/internal/store/worker_input_race_test.go b/services/core/internal/store/worker_input_race_test.go index c12653dc4..6ca03ded1 100644 --- a/services/core/internal/store/worker_input_race_test.go +++ b/services/core/internal/store/worker_input_race_test.go @@ -50,7 +50,7 @@ func TestWorkerInputReadSkipsConcurrentlyCancelledCandidate(t *testing.T) { mutated <- h.s.CommitLegacyDeletion(t.Context(), h.tenant, candidateSession) return } - _, err := h.s.SubmitInputs(t.Context(), h.tenant, candidateSession, "cancel", []sessions.Input{{Kind: "cancel", Payload: json.RawMessage(`{}`)}}) + _, err := store.SubmitInputs(t.Context(), h.s, h.tenant, candidateSession, "cancel", []sessions.Input{{Kind: "cancel", Payload: json.RawMessage(`{}`)}}) mutated <- err }} instrumented, err := pgxpool.NewWithConfig(t.Context(), cfg) diff --git a/services/core/internal/store/worker_preparation_failure_test.go b/services/core/internal/store/worker_preparation_failure_test.go index 93b293053..b7c82413e 100644 --- a/services/core/internal/store/worker_preparation_failure_test.go +++ b/services/core/internal/store/worker_preparation_failure_test.go @@ -16,7 +16,7 @@ func TestWorkerSettlesConfirmedPreparationFailureAndAcceptsNewInput(t *testing.T h := newDispatchHarnessForSession(t, []byte(`{"agent":{"model":"test-model"},"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`), false) enableWorkerEnvironment(t, h) frames := workerFrames(t, h) - pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "first", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) + pending, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "first", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) if err != nil { t.Fatal(err) } @@ -31,7 +31,7 @@ func TestWorkerSettlesConfirmedPreparationFailureAndAcceptsNewInput(t *testing.T h.write(prepare.ID, proto.TypePreparationStatus, proto.PreparationStatusPayload{State: "rejected", Operation: proto.TypeExecutionPrepare, ErrorCode: code}) } awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "failed reservation settlement", func() bool { - current, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) + current, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) return err == nil && current.State == sessions.EnvironmentInputFailed }) session, err := store.SessionAdapter(h.s).GetSession(t.Context(), h.tenant, h.session.ID) @@ -59,7 +59,7 @@ func TestWorkerSettlesConfirmedPreparationFailureAndAcceptsNewInput(t *testing.T t.Fatal("failed input retried", frame.Type) case <-time.After(1200 * time.Millisecond): } - next, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "next", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"next"}`)}}) + next, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "next", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"next"}`)}}) if err != nil { t.Fatal("new input remained blocked", err) } @@ -71,7 +71,7 @@ func TestWorkerSettlesConfirmedPreparationFailureAndAcceptsNewInput(t *testing.T if startFrame.DecodePayload(&start) != nil || start.RunID == "" || inputTextForTest(t, start.Input) != "next" { t.Fatal("new input was not admitted") } - current, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, next.ID) + current, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, next.ID) if err != nil || current.State != sessions.EnvironmentInputAdmitted { t.Fatal("new input state", err) } @@ -99,7 +99,7 @@ func TestWorkerRetriesUncertainPreparationFailure(t *testing.T) { h := newDispatchHarnessForSession(t, []byte(`{"agent":{"model":"test-model"},"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`), false) enableWorkerEnvironment(t, h) frames := workerFrames(t, h) - pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "retry", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"retry"}`)}}) + pending, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "retry", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"retry"}`)}}) if err != nil { t.Fatal(err) } @@ -116,11 +116,11 @@ func TestWorkerRetriesUncertainPreparationFailure(t *testing.T) { nextWorkerFrame(t, frames, proto.TypeExecutionRelease) } nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) - current, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) + current, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, pending.ID) if err != nil || current.State != sessions.EnvironmentInputPending || !current.Deadline.Equal(pending.Deadline) { t.Fatal("transient failure settled or extended input", err) } - if _, err := h.s.CancelEnvironmentInput(t.Context(), h.tenant, h.session.ID, pending.ID); err != nil { + if _, err := store.CancelEnvironmentInput(t.Context(), h.s, h.tenant, h.session.ID, pending.ID); err != nil { t.Fatal(err) } }) @@ -131,17 +131,17 @@ func TestWorkerPreparationRejectionPreservesCancellationAndNewerInput(t *testing h := newDispatchHarnessForSession(t, []byte(`{"agent":{"model":"test-model"},"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`), false) enableWorkerEnvironment(t, h) frames := workerFrames(t, h) - first, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "first", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) + first, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "first", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) if err != nil { t.Fatal(err) } _, stop := startEnvironmentExpiryWorker(t, h.db, h.d) defer stop() old := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) - if _, err := h.s.CancelEnvironmentInput(t.Context(), h.tenant, h.session.ID, first.ID); err != nil { + if _, err := store.CancelEnvironmentInput(t.Context(), h.s, h.tenant, h.session.ID, first.ID); err != nil { t.Fatal(err) } - next, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "next", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"next"}`)}}) + next, err := store.SessionService(t, h.s).ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "next", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"next"}`)}}) if err != nil { t.Fatal(err) } @@ -149,7 +149,7 @@ func TestWorkerPreparationRejectionPreservesCancellationAndNewerInput(t *testing h.write(old.ID, proto.TypePreparationStatus, rejection) prepare := nextWorkerFrame(t, frames, proto.TypeExecutionPrepare) for id, state := range map[string]string{first.ID: sessions.EnvironmentInputCancelled, next.ID: sessions.EnvironmentInputPending} { - current, err := h.s.GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, id) + current, err := store.SessionAdapter(h.s).GetEnvironmentInputReservation(t.Context(), h.tenant, h.session.ID, id) if err != nil || current.State != state { t.Fatal("late rejection changed cancellation or newer input", err) } diff --git a/services/core/tests/fixtures/main.go b/services/core/tests/fixtures/main.go index 936afddfd..12f1e8159 100644 --- a/services/core/tests/fixtures/main.go +++ b/services/core/tests/fixtures/main.go @@ -8,8 +8,9 @@ import ( "os" "strings" + "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" ) @@ -52,12 +53,16 @@ func seed() error { return err } defer pool.Close() - s := store.New(pool) + service, err := sessions.NewService(sessionpg.New(pgunit.NewPool(pool), nil)) + if err != nil { + return err + } for _, status := range []string{sessions.TurnCompleted, sessions.TurnFailed, sessions.TurnCancelled, sessions.TurnInProgress} { - receipt, err := s.SubmitMessage(ctx, f.Tenant, f.Session, uuid.NewString(), json.RawMessage(`{"text":"recovery fixture"}`)) + receipts, err := service.SubmitInputs(ctx, f.Tenant, f.Session, uuid.NewString(), []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"recovery fixture"}`)}}) if err != nil { return err } + receipt := receipts[0] if err = transitionTurn(ctx, pool, f.Tenant, f.Session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { return err }