From 400b60b9f5ce505d160291b343be7e4429ee8fe8 Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 30 Sep 2026 20:21:51 +0000 Subject: [PATCH] Share Session transaction procedures through sessions and sessionpg Cancellation, Environment failure and termination, the input-activity bracket, the admission gates and dedicated device creation become sessions procedures over the transaction interfaces they declare, with pure decisions and explicit post-update loads. sessionpg adds LockSession, WithSession (lock, apply, prune on a pgunit.Transactor), the SessionTx binding from BindSession, which writes only its own Session's Environment, LoadEnvironmentInput and the row helpers. store's private helpers become direct calls; lock order, error precedence, event order and rollback are unchanged. --- services/core/IMPLEMENTATION.md | 4 +- .../db/queries/environment_connections.sql | 12 +- .../db/queries/environment_input_expiry.sql | 2 +- .../db/queries/local_environment_devices.sql | 2 +- .../db/sqlc/environment_connections.sql.go | 39 ++- .../db/sqlc/environment_input_expiry.sql.go | 5 +- .../db/sqlc/local_environment_devices.sql.go | 4 +- .../persistence/postgres/pgunit/pgunit.go | 11 + .../persistence/postgres/sessionpg/binding.go | 192 ++++++++++++ .../persistence/postgres/sessionpg/doc.go | 16 +- .../postgres/sessionpg/environment.go | 34 ++ .../persistence/postgres/sessionpg/rows.go | 120 +++++++ .../persistence/postgres/sessionpg/session.go | 49 +++ .../postgres/sessionpg/session_test.go | 296 ++++++++++++++++++ services/core/internal/sessions/admission.go | 71 +++++ .../core/internal/sessions/admission_test.go | 43 +++ .../core/internal/sessions/cancellation.go | 105 +++++++ .../internal/sessions/cancellation_test.go | 87 +++++ services/core/internal/sessions/doc.go | 4 + .../internal/sessions/environment_device.go | 31 ++ .../sessions/environment_device_test.go | 22 ++ .../sessions/environment_termination.go | 165 ++++++++++ .../sessions/environment_termination_test.go | 133 ++++++++ .../core/internal/sessions/input_activity.go | 125 ++++++++ .../internal/sessions/input_activity_test.go | 103 ++++++ .../core/internal/sessions/transaction.go | 44 +++ .../internal/sessions/transaction_test.go | 193 ++++++++++++ .../internal/store/admin_session_archive.go | 8 +- services/core/internal/store/admin_summary.go | 3 +- .../core/internal/store/artifact_capture.go | 13 +- .../core/internal/store/artifact_lifecycle.go | 3 +- services/core/internal/store/devices.go | 21 +- .../internal/store/environment_connections.go | 37 +-- .../internal/store/environment_file_writes.go | 14 +- .../store/environment_initialization.go | 7 +- .../store/environment_input_activity.go | 84 +---- .../store/environment_input_expiry.go | 3 +- .../core/internal/store/environment_inputs.go | 27 +- .../core/internal/store/environment_setup.go | 3 +- services/core/internal/store/environments.go | 26 +- services/core/internal/store/execution.go | 7 - .../core/internal/store/function_calls.go | 7 +- .../core/internal/store/function_results.go | 3 +- .../core/internal/store/function_state.go | 4 +- services/core/internal/store/initial_files.go | 3 +- .../store/local_environment_devices.go | 38 +-- .../store/runtime_allocation_state.go | 9 +- .../internal/store/runtime_allocations.go | 8 +- .../core/internal/store/runtime_enrollment.go | 3 +- .../store/runtime_environment_terminal.go | 97 ------ .../internal/store/runtime_lifecycle_nodes.go | 3 +- .../core/internal/store/runtime_suspension.go | 13 +- .../runtime_suspension_concurrency_test.go | 3 +- services/core/internal/store/scheduling.go | 23 +- .../core/internal/store/session_artifacts.go | 9 +- .../store/session_creation_identity.go | 3 +- .../core/internal/store/session_creator.go | 19 -- .../core/internal/store/session_deletion.go | 67 +--- .../internal/store/session_deletion_test.go | 7 +- .../internal/store/session_diagnostics.go | 7 +- .../store/session_environment_snapshot.go | 3 +- .../internal/store/session_initial_input.go | 2 +- .../core/internal/store/session_metadata.go | 3 +- .../internal/store/session_transaction.go | 36 +-- services/core/internal/store/sessions.go | 27 +- .../internal/store/subagent_identities.go | 3 +- .../core/internal/store/turn_completion.go | 6 +- services/core/internal/store/turn_events.go | 5 +- services/core/internal/store/turn_inputs.go | 18 +- services/core/internal/store/turn_reads.go | 3 +- services/core/internal/store/turns.go | 36 +-- 71 files changed, 2080 insertions(+), 556 deletions(-) create mode 100644 services/core/internal/persistence/postgres/sessionpg/binding.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/environment.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/rows.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/session.go create mode 100644 services/core/internal/persistence/postgres/sessionpg/session_test.go create mode 100644 services/core/internal/sessions/admission.go create mode 100644 services/core/internal/sessions/admission_test.go create mode 100644 services/core/internal/sessions/cancellation.go create mode 100644 services/core/internal/sessions/cancellation_test.go create mode 100644 services/core/internal/sessions/environment_device.go create mode 100644 services/core/internal/sessions/environment_device_test.go create mode 100644 services/core/internal/sessions/environment_termination.go create mode 100644 services/core/internal/sessions/environment_termination_test.go create mode 100644 services/core/internal/sessions/input_activity.go create mode 100644 services/core/internal/sessions/input_activity_test.go create mode 100644 services/core/internal/sessions/transaction.go create mode 100644 services/core/internal/sessions/transaction_test.go delete mode 100644 services/core/internal/store/runtime_environment_terminal.go delete mode 100644 services/core/internal/store/session_creator.go diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 1a1463257..54b1fd32d 100644 --- a/services/core/IMPLEMENTATION.md +++ b/services/core/IMPLEMENTATION.md @@ -16,7 +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. `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 and output indexes, writes the journal, Items, Turn usage and Artifact settlement, and prunes the journal. It decides nothing. +`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, and `CheckComputeAdmission`, `CheckFileWriteGate` and `CheckInputStart` gate new work. `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 and output indexes, writes the journal, Items, Turn usage and Artifact settlement, and prunes the 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 journal. `BindSession` binds one Session to a transaction's queries as a `SessionTx`, which implements the procedures' interfaces. + +Other adapters call only two `sessionpg` participants, on their own transaction's queries: `LockSession`, which orders a write against the Session's Turns, and `BindSession`, which runs a Session procedure, such as `CreateEnvironmentDevice`, in that transaction. The other `sessionpg` exports serve `store` until the Session operations leave it. `deploymentpg` may import `sessionpg`, and `sessionpg` never imports `deploymentpg`; deployment's placement fact-loading and reservation code sits below both, so either may call it and it imports neither. Likewise `deployment` may compose the `sessions` procedures, and `sessions` never imports `deployment`. `store` is transitional. `store.New` builds a pooled Store, and `store.NewExecution(s, lease)` builds the execution writer on a lease it borrows. An execution-only operation on a pooled Store fails with `store.ErrExecutionAuthority`. New adapters do not copy that check: their execution repositories require a `*pgunit.Lease` at construction, their public repositories expose no execution operation, and the check goes away with `store`. diff --git a/services/core/internal/db/queries/environment_connections.sql b/services/core/internal/db/queries/environment_connections.sql index d9461f24b..3f0d8a2c5 100644 --- a/services/core/internal/db/queries/environment_connections.sql +++ b/services/core/internal/db/queries/environment_connections.sql @@ -13,9 +13,15 @@ UPDATE environment_connections SET revision = $2 WHERE environment_id = $1; UPDATE environments SET status = $2 WHERE id = $1; -- name: RecordEnvironmentFailure :one -UPDATE environments SET status = 'failed', failure_reason = $2, failure_detail = $3, failed_at = clock_timestamp() -WHERE id = $1 AND status NOT IN ('failed', 'expired') -RETURNING failed_at; +UPDATE environments e SET status = 'failed', failure_reason = sqlc.arg(failure_reason), failure_detail = sqlc.arg(failure_detail), failed_at = clock_timestamp() +WHERE e.id = sqlc.arg(id) AND e.session_id = sqlc.arg(session_id) AND e.status NOT IN ('failed', 'expired') +AND EXISTS (SELECT 1 FROM sessions s WHERE s.id = e.session_id AND s.tenant_id = sqlc.arg(tenant_id)) +RETURNING e.failed_at; + +-- name: ExpireSessionEnvironment :execrows +UPDATE environments e SET status = 'expired' +WHERE e.id = sqlc.arg(id) AND e.session_id = sqlc.arg(session_id) +AND EXISTS (SELECT 1 FROM sessions s WHERE s.id = e.session_id AND s.tenant_id = sqlc.arg(tenant_id)); -- name: DeleteEnvironmentConnection :exec DELETE FROM environment_connections WHERE environment_id = $1; diff --git a/services/core/internal/db/queries/environment_input_expiry.sql b/services/core/internal/db/queries/environment_input_expiry.sql index afa549fd9..c4e23c536 100644 --- a/services/core/internal/db/queries/environment_input_expiry.sql +++ b/services/core/internal/db/queries/environment_input_expiry.sql @@ -1,5 +1,5 @@ -- name: ListDueEnvironmentInputs :many -SELECT r.id, r.session_id +SELECT r.id, r.session_id, s.tenant_id FROM environment_input_reservations r JOIN sessions s ON s.id = r.session_id WHERE r.state = 'pending' diff --git a/services/core/internal/db/queries/local_environment_devices.sql b/services/core/internal/db/queries/local_environment_devices.sql index 29b917620..75d50bfc5 100644 --- a/services/core/internal/db/queries/local_environment_devices.sql +++ b/services/core/internal/db/queries/local_environment_devices.sql @@ -2,7 +2,7 @@ INSERT INTO devices (id, tenant_id, name, credential_hash, environment_id) SELECT sqlc.arg(id), s.tenant_id, sqlc.arg(name), sqlc.arg(credential_hash), e.id FROM environments e JOIN sessions s ON s.id = e.session_id -WHERE s.tenant_id = sqlc.arg(tenant_id) AND e.id = sqlc.arg(environment_id) +WHERE s.tenant_id = sqlc.arg(tenant_id) AND s.id = sqlc.arg(session_id) AND e.id = sqlc.arg(environment_id) AND s.deleted_at IS NULL AND s.configuration->'environment'->>'type' = 'openai_hosted' ON CONFLICT (environment_id) DO NOTHING RETURNING id; diff --git a/services/core/internal/db/sqlc/environment_connections.sql.go b/services/core/internal/db/sqlc/environment_connections.sql.go index 57ddde24f..96d60f5a8 100644 --- a/services/core/internal/db/sqlc/environment_connections.sql.go +++ b/services/core/internal/db/sqlc/environment_connections.sql.go @@ -34,6 +34,26 @@ func (q *Queries) DeleteEnvironmentConnection(ctx context.Context, environmentID return err } +const expireSessionEnvironment = `-- name: ExpireSessionEnvironment :execrows +UPDATE environments e SET status = 'expired' +WHERE e.id = $1 AND e.session_id = $2 +AND EXISTS (SELECT 1 FROM sessions s WHERE s.id = e.session_id AND s.tenant_id = $3) +` + +type ExpireSessionEnvironmentParams struct { + ID pgtype.UUID `json:"id"` + SessionID pgtype.UUID `json:"session_id"` + TenantID pgtype.UUID `json:"tenant_id"` +} + +func (q *Queries) ExpireSessionEnvironment(ctx context.Context, arg ExpireSessionEnvironmentParams) (int64, error) { + result, err := q.db.Exec(ctx, expireSessionEnvironment, arg.ID, arg.SessionID, arg.TenantID) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + const getEnvironmentConnection = `-- name: GetEnvironmentConnection :one SELECT environment_id, generation, revision FROM environment_connections WHERE environment_id = $1 ` @@ -82,19 +102,28 @@ func (q *Queries) ListEnvironmentConnections(ctx context.Context, id pgtype.UUID } const recordEnvironmentFailure = `-- name: RecordEnvironmentFailure :one -UPDATE environments SET status = 'failed', failure_reason = $2, failure_detail = $3, failed_at = clock_timestamp() -WHERE id = $1 AND status NOT IN ('failed', 'expired') -RETURNING failed_at +UPDATE environments e SET status = 'failed', failure_reason = $1, failure_detail = $2, failed_at = clock_timestamp() +WHERE e.id = $3 AND e.session_id = $4 AND e.status NOT IN ('failed', 'expired') +AND EXISTS (SELECT 1 FROM sessions s WHERE s.id = e.session_id AND s.tenant_id = $5) +RETURNING e.failed_at ` type RecordEnvironmentFailureParams struct { - ID pgtype.UUID `json:"id"` FailureReason pgtype.Text `json:"failure_reason"` FailureDetail []byte `json:"failure_detail"` + ID pgtype.UUID `json:"id"` + SessionID pgtype.UUID `json:"session_id"` + TenantID pgtype.UUID `json:"tenant_id"` } func (q *Queries) RecordEnvironmentFailure(ctx context.Context, arg RecordEnvironmentFailureParams) (pgtype.Timestamptz, error) { - row := q.db.QueryRow(ctx, recordEnvironmentFailure, arg.ID, arg.FailureReason, arg.FailureDetail) + row := q.db.QueryRow(ctx, recordEnvironmentFailure, + arg.FailureReason, + arg.FailureDetail, + arg.ID, + arg.SessionID, + arg.TenantID, + ) var failed_at pgtype.Timestamptz err := row.Scan(&failed_at) return failed_at, err diff --git a/services/core/internal/db/sqlc/environment_input_expiry.sql.go b/services/core/internal/db/sqlc/environment_input_expiry.sql.go index a1fbacf18..410729fa9 100644 --- a/services/core/internal/db/sqlc/environment_input_expiry.sql.go +++ b/services/core/internal/db/sqlc/environment_input_expiry.sql.go @@ -12,7 +12,7 @@ import ( ) const listDueEnvironmentInputs = `-- name: ListDueEnvironmentInputs :many -SELECT r.id, r.session_id +SELECT r.id, r.session_id, s.tenant_id FROM environment_input_reservations r JOIN sessions s ON s.id = r.session_id WHERE r.state = 'pending' @@ -26,6 +26,7 @@ FOR UPDATE OF s SKIP LOCKED type ListDueEnvironmentInputsRow struct { ID pgtype.UUID `json:"id"` SessionID pgtype.UUID `json:"session_id"` + TenantID pgtype.UUID `json:"tenant_id"` } func (q *Queries) ListDueEnvironmentInputs(ctx context.Context) ([]ListDueEnvironmentInputsRow, error) { @@ -37,7 +38,7 @@ func (q *Queries) ListDueEnvironmentInputs(ctx context.Context) ([]ListDueEnviro items := []ListDueEnvironmentInputsRow{} for rows.Next() { var i ListDueEnvironmentInputsRow - if err := rows.Scan(&i.ID, &i.SessionID); err != nil { + if err := rows.Scan(&i.ID, &i.SessionID, &i.TenantID); err != nil { return nil, err } items = append(items, i) diff --git a/services/core/internal/db/sqlc/local_environment_devices.sql.go b/services/core/internal/db/sqlc/local_environment_devices.sql.go index 51c15c33b..72c89655c 100644 --- a/services/core/internal/db/sqlc/local_environment_devices.sql.go +++ b/services/core/internal/db/sqlc/local_environment_devices.sql.go @@ -15,7 +15,7 @@ const createEnvironmentDevice = `-- name: CreateEnvironmentDevice :one INSERT INTO devices (id, tenant_id, name, credential_hash, environment_id) SELECT $1, s.tenant_id, $2, $3, e.id FROM environments e JOIN sessions s ON s.id = e.session_id -WHERE s.tenant_id = $4 AND e.id = $5 +WHERE s.tenant_id = $4 AND s.id = $5 AND e.id = $6 AND s.deleted_at IS NULL AND s.configuration->'environment'->>'type' = 'openai_hosted' ON CONFLICT (environment_id) DO NOTHING RETURNING id @@ -26,6 +26,7 @@ type CreateEnvironmentDeviceParams struct { Name string `json:"name"` CredentialHash pgtype.Text `json:"credential_hash"` TenantID pgtype.UUID `json:"tenant_id"` + SessionID pgtype.UUID `json:"session_id"` EnvironmentID pgtype.UUID `json:"environment_id"` } @@ -35,6 +36,7 @@ func (q *Queries) CreateEnvironmentDevice(ctx context.Context, arg CreateEnviron arg.Name, arg.CredentialHash, arg.TenantID, + arg.SessionID, arg.EnvironmentID, ) var id pgtype.UUID diff --git a/services/core/internal/persistence/postgres/pgunit/pgunit.go b/services/core/internal/persistence/postgres/pgunit/pgunit.go index 1bd064a11..f7032c4cb 100644 --- a/services/core/internal/persistence/postgres/pgunit/pgunit.go +++ b/services/core/internal/persistence/postgres/pgunit/pgunit.go @@ -19,6 +19,17 @@ var ( snapshot = pgx.TxOptions{IsoLevel: pgx.RepeatableRead, AccessMode: pgx.ReadOnly} ) +// Transactor runs read committed transactions: a Pool on pooled connections, +// or the Lease on the execution connection. +type Transactor interface { + Transaction(ctx context.Context, apply func(context.Context, pgx.Tx) error) error +} + +var ( + _ Transactor = (*Pool)(nil) + _ Transactor = (*Lease)(nil) +) + // Pool runs transactions on pooled connections. It grants no execution authority. type Pool struct{ pool *pgxpool.Pool } diff --git a/services/core/internal/persistence/postgres/sessionpg/binding.go b/services/core/internal/persistence/postgres/sessionpg/binding.go new file mode 100644 index 000000000..79c14f4c0 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/binding.go @@ -0,0 +1,192 @@ +package sessionpg + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// SessionTx is one Session inside its caller's transaction. It implements the +// transaction interfaces of the sessions procedures by loading the facts they +// read and applying what they decide, and decides nothing itself. +// +// The caller guarantees what the procedures assume: the queries are bound to +// the caller's open transaction, the tenant owns the Session, the transaction +// already holds the Session lock (LockSession), and the operation that owns the +// transaction supplies its authority, such as the execution lease. A SessionTx +// never begins or commits a transaction, never chooses a pool or the lease, and +// is not used after the transaction callback it was bound in returns. +type SessionTx struct { + q *sqlc.Queries + tenant, session pgtype.UUID +} + +var ( + _ sessions.EnvironmentTerminationTx = (*SessionTx)(nil) + _ sessions.InputStartTx = (*SessionTx)(nil) + _ sessions.ComputeAdmissionTx = (*SessionTx)(nil) + _ sessions.EnvironmentDeviceTx = (*SessionTx)(nil) +) + +// BindSession binds the tenant's Session to the caller's transaction-bound +// queries. +func BindSession(q *sqlc.Queries, tenant, session pgtype.UUID) *SessionTx { + return &SessionTx{q: q, tenant: tenant, session: session} +} + +func (t *SessionTx) LoadUsage(ctx context.Context) (json.RawMessage, error) { + return LoadUsage(ctx, t.q, t.session) +} + +func (t *SessionTx) AppendChanges(ctx context.Context, changes ...sessions.SessionChange) error { + return AppendChanges(ctx, t.q, t.session, changes...) +} + +func (t *SessionTx) LoadActiveTurn(ctx context.Context) (sessions.Turn, bool, error) { + row, err := t.q.GetActiveTurn(ctx, t.session) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.Turn{}, false, nil + } + if err != nil { + return sessions.Turn{}, false, err + } + return TurnFromRow(row), true, nil +} + +func (t *SessionTx) LoadTurn(ctx context.Context, turn string) (sessions.Turn, error) { + id, err := parseID(turn) + if err != nil { + return sessions.Turn{}, err + } + row, err := t.q.SessionEventTurn(ctx, sqlc.SessionEventTurnParams{SessionID: t.session, ID: id}) + if err != nil { + return sessions.Turn{}, err + } + return TurnFromRow(row), nil +} + +func (t *SessionTx) RequestTurnCancel(ctx context.Context, turn string) error { + id, err := parseID(turn) + if err != nil { + return err + } + return t.q.RequestTurnCancel(ctx, sqlc.RequestTurnCancelParams{ID: id, SessionID: t.session}) +} + +func (t *SessionTx) LoadEnding(ctx context.Context, turn string) (sessions.Ending, error) { + id, err := parseID(turn) + if err != nil { + return sessions.Ending{}, err + } + return LoadEnding(ctx, t.q, t.session, id) +} + +func (t *SessionTx) ApplyTurnEnd(ctx context.Context, turn string, end sessions.TurnEnd) error { + id, err := parseID(turn) + if err != nil { + return err + } + return ApplyTurnEnd(ctx, t.q, t.session, id, end) +} + +func (t *SessionTx) CancelPendingInput(ctx context.Context) error { + return t.q.CancelSessionEnvironmentInput(ctx, t.session) +} + +func (t *SessionTx) FailPendingInput(ctx context.Context) error { + return t.q.FailSessionEnvironmentInput(ctx, t.session) +} + +func (t *SessionTx) LoadEnvironmentInput(ctx context.Context) (*sessions.EnvironmentInputState, error) { + return LoadEnvironmentInput(ctx, t.q, t.session) +} + +func (t *SessionTx) LoadEnvironment(ctx context.Context) (sessions.Environment, error) { + row, err := t.q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: t.tenant, ID: t.session}) + return EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, err) +} + +// RecordEnvironmentFailure fails the Environment only while it is the +// Session's and still live; otherwise it changes nothing and returns +// pgx.ErrNoRows. +func (t *SessionTx) RecordEnvironmentFailure(ctx context.Context, environment, reason string, detail *sessions.ProvisioningFailureDetail) (time.Time, error) { + id, err := parseID(environment) + if err != nil { + return time.Time{}, err + } + var rawDetail []byte + if detail != nil { + if rawDetail, err = json.Marshal(detail); err != nil { + return time.Time{}, err + } + } + failedAt, err := t.q.RecordEnvironmentFailure(ctx, sqlc.RecordEnvironmentFailureParams{ + ID: id, SessionID: t.session, TenantID: t.tenant, + FailureReason: pgtype.Text{String: reason, Valid: true}, FailureDetail: rawDetail, + }) + return failedAt.Time, err +} + +// ExpireEnvironment expires the Environment only while it is the Session's; +// another Environment is sessions.ErrNotFound and changes nothing. +func (t *SessionTx) ExpireEnvironment(ctx context.Context, environment string) error { + id, err := parseID(environment) + if err != nil { + return err + } + count, err := t.q.ExpireSessionEnvironment(ctx, sqlc.ExpireSessionEnvironmentParams{ID: id, SessionID: t.session, TenantID: t.tenant}) + if err == nil && count != 1 { + return sessions.ErrNotFound + } + return err +} + +func (t *SessionTx) LoadBoundDevice(ctx context.Context) (bool, error) { + _, err := t.q.GetSessionDevice(ctx, sqlc.GetSessionDeviceParams{TenantID: t.tenant, ID: t.session}) + if errors.Is(err, pgx.ErrNoRows) { + return false, nil + } + return err == nil, err +} + +// InsertEnvironmentDevice inserts the device only for the Session's own hosted +// Environment. The insert takes no row when the Environment already has a +// device, including one a concurrent transaction inserted first, or cannot +// take one: that is sessions.ErrDeviceBindingConflict. +func (t *SessionTx) InsertEnvironmentDevice(ctx context.Context, device sessions.ExecutionDevice, credentialHash string) error { + id, err := parseID(device.ID) + if err != nil { + return err + } + environment, err := parseID(device.EnvironmentID) + if err != nil { + return err + } + inserted, err := t.q.CreateEnvironmentDevice(ctx, sqlc.CreateEnvironmentDeviceParams{ + ID: id, TenantID: t.tenant, SessionID: t.session, EnvironmentID: environment, + Name: device.Name, CredentialHash: pgtype.Text{String: credentialHash, Valid: true}, + }) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.ErrDeviceBindingConflict + } + if err != nil { + return err + } + _, err = t.q.BindSessionDevice(ctx, sqlc.BindSessionDeviceParams{TenantID: t.tenant, ID: t.session, ID_2: inserted}) + return err +} + +func (t *SessionTx) LoadComputeSuspension(ctx context.Context) (bool, error) { + return t.q.RuntimeComputeBlocksAdmission(ctx, t.session) +} + +func (t *SessionTx) LoadPendingFileWrite(ctx context.Context) (bool, error) { + return t.q.EnvironmentFileWriteBlocksSession(ctx, t.session) +} diff --git a/services/core/internal/persistence/postgres/sessionpg/doc.go b/services/core/internal/persistence/postgres/sessionpg/doc.go index 9f0c83f12..4709ade9d 100644 --- a/services/core/internal/persistence/postgres/sessionpg/doc.go +++ b/services/core/internal/persistence/postgres/sessionpg/doc.go @@ -1,10 +1,12 @@ -// Package sessionpg writes the Session changes that the sessions and items -// packages decide to PostgreSQL. Its functions run inside the caller's Session -// transaction, under the Session lock, on the transaction-bound queries: they -// load the facts a decision reads, then apply the decision, allocating event -// sequence positions, event IDs and Item positions, writing the public change -// journal, Items, Turn usage and Artifacts, and pruning the journal. They decide -// nothing. +// Package sessionpg stores Sessions in PostgreSQL and decides nothing. +// WithSession runs a Session transaction: it locks the Session, runs the +// operation and prunes the journal. Inside a Session transaction, on its +// queries, the participants load the facts the sessions and items decisions +// read and apply what they decide: they lock the Session, allocate event +// sequence positions, event IDs, Item positions and output indexes, and write +// the public change journal, Items, Turn usage, Artifact settlement, +// Environment state, input reservations and devices. SessionTx binds a Session +// to its caller's transaction for the sessions procedures. package sessionpg import ( diff --git a/services/core/internal/persistence/postgres/sessionpg/environment.go b/services/core/internal/persistence/postgres/sessionpg/environment.go new file mode 100644 index 000000000..fd6667bfe --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/environment.go @@ -0,0 +1,34 @@ +package sessionpg + +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/sessions" +) + +// LoadEnvironmentInput reads the Session's latest Environment input +// reservation that no Turn has admitted or superseded, with the facts of its +// Environment, and nil when there is none. sessions.InputActivity projects it. +func LoadEnvironmentInput(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) (*sessions.EnvironmentInputState, error) { + row, err := q.GetEnvironmentInputActivity(ctx, session) + if errors.Is(err, pgx.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, err + } + state := &sessions.EnvironmentInputState{ + State: row.State, Initial: row.IsInitial, CreatedAt: row.CreatedAt.Time, FailureCode: row.FailureCode.String, + EnvironmentID: uuid.UUID(row.EnvironmentID.Bytes).String(), EnvironmentType: row.EnvironmentType, EnvironmentStatus: row.ConnectionStatus, + } + if row.SettledAt.Valid { + state.SettledAt = row.SettledAt.Time + } + return state, nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/rows.go b/services/core/internal/persistence/postgres/sessionpg/rows.go new file mode 100644 index 000000000..f4b885714 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/rows.go @@ -0,0 +1,120 @@ +package sessionpg + +import ( + "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/identity" + "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/sessions" +) + +// parseID parses an internal identifier; a malformed one is +// sessions.ErrInvalidInput. +func parseID(value string) (pgtype.UUID, error) { + id, err := pgunit.ParseID(value) + if err != nil { + return pgtype.UUID{}, fmt.Errorf("%w: %w", sessions.ErrInvalidInput, err) + } + return id, nil +} + +// TurnLookup parses the internal identifiers of a tenant's Turn; a malformed +// one is sessions.ErrInvalidInput. +func TurnLookup(tenantID, sessionID, turnID string) (sqlc.GetTurnParams, error) { + var p sqlc.GetTurnParams + var err error + if p.TenantID, err = parseID(tenantID); err != nil { + return p, err + } + if p.SessionID, err = parseID(sessionID); err != nil { + return p, err + } + p.ID, err = parseID(turnID) + return p, err +} + +// DeviceLookup parses a tenant and the internal identifier of one of its +// resources, such as a device, Session or Environment; a malformed one is +// sessions.ErrInvalidInput. +func DeviceLookup(tenantID, id string) (sqlc.GetDeviceParams, error) { + var p sqlc.GetDeviceParams + var err error + if p.TenantID, err = parseID(tenantID); err != nil { + return p, err + } + p.ID, err = parseID(id) + return p, err +} + +// TurnFromRow maps a stored Turn. +func TurnFromRow(row sqlc.Turn) sessions.Turn { + return sessions.Turn{ + ID: uuid.UUID(row.ID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), Status: row.Status, + CreatedAt: row.CreatedAt.Time, StartedAt: row.StartedAt.Time, CompletedAt: row.CompletedAt.Time, + CancelRequestedAt: row.CancelRequestedAt.Time, Outcome: json.RawMessage(row.Outcome), Usage: json.RawMessage(row.TokenUsage), + ArtifactCaptureStarted: row.ArtifactCaptureStarted, + } +} + +// SessionFromRow maps a stored Session with its creator, normalized +// configuration and metadata. +func SessionFromRow(row sqlc.Session) (sessions.Session, error) { + session := sessions.Session{ID: uuid.UUID(row.ID.Bytes).String(), TenantID: uuid.UUID(row.TenantID.Bytes).String(), Engine: row.Engine, CreatedAt: row.CreatedAt.Time, RequiredActions: []v1.FunctionCallAction{}} + creator, err := sessionCreator(row.CreatorKind, row.CreatorID) + if err != nil { + return sessions.Session{}, err + } + session.Creator = creator + configuration, err := jsonobject.Normalize(row.Configuration) + if err != nil { + return sessions.Session{}, fmt.Errorf("decode session configuration: %w: %w", sessions.ErrInvalidInput, err) + } + session.Configuration = configuration + if err := json.Unmarshal(row.Metadata, &session.Metadata); err != nil { + return sessions.Session{}, fmt.Errorf("decode session metadata: %w", err) + } + return session, nil +} + +// sessionCreator maps the stored creator: nil when none is recorded, and an +// error for a partial or invalid one. +func sessionCreator(kind, id pgtype.Text) (*identity.Subject, error) { + if !kind.Valid && !id.Valid { + return nil, nil + } + creator := identity.Subject{Kind: kind.String, ID: id.String} + if !kind.Valid || !id.Valid || creator.Validate() != nil { + return nil, fmt.Errorf("invalid stored Session creator") + } + return &creator, nil +} + +// EnvironmentFromRow maps a stored Environment of the tenant with its +// normalized configuration snapshot, from the result of the query that read +// it: no rows is sessions.ErrNotFound. +func EnvironmentFromRow(row sqlc.Environment, tenant pgtype.UUID, configuration []byte, err error) (sessions.Environment, error) { + if errors.Is(err, pgx.ErrNoRows) { + return sessions.Environment{}, sessions.ErrNotFound + } + if err != nil { + return sessions.Environment{}, fmt.Errorf("get environment: %w", err) + } + configuration, err = jsonobject.Normalize(configuration) + if err != nil { + return sessions.Environment{}, fmt.Errorf("decode environment configuration: %w: %w", sessions.ErrInvalidInput, err) + } + return sessions.Environment{ + ID: uuid.UUID(row.ID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), + TenantID: uuid.UUID(tenant.Bytes).String(), Status: row.Status, + Initialization: row.Initialization, CreatedAt: row.CreatedAt.Time, Configuration: configuration, + }, nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/session.go b/services/core/internal/persistence/postgres/sessionpg/session.go new file mode 100644 index 000000000..162dabe26 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/session.go @@ -0,0 +1,49 @@ +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/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// LockSession locks the tenant's Session row, including a publicly deleted +// one, and returns what the lock shows. Every Session write takes this lock +// first, which orders input admission against completion and cancellation +// across service processes. A missing Session is sessions.ErrNotFound. +func LockSession(ctx context.Context, q *sqlc.Queries, tenant, session pgtype.UUID) (sessions.LockedSession, error) { + row, err := q.LockSession(ctx, sqlc.LockSessionParams{TenantID: tenant, ID: session}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.LockedSession{}, sessions.ErrNotFound + } + if err != nil { + return sessions.LockedSession{}, err + } + return sessions.LockedSession{Deleted: row.DeletedAt.Valid}, nil +} + +// WithSession runs one Session transaction on runner: it locks the tenant's +// Session, runs apply on the transaction's queries with what the lock shows, +// then prunes the Session's journal, and commits only when all of them +// succeed. apply applies the public view itself, with +// sessions.LockedSession.Public. It is the transaction runner of Session +// operations; other adapters lock the Session inside their own transaction +// with LockSession. +func WithSession(ctx context.Context, runner pgunit.Transactor, tenant, session pgtype.UUID, apply func(context.Context, *sqlc.Queries, sessions.LockedSession) error) error { + return runner.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + locked, err := LockSession(ctx, q, tenant, session) + if err != nil { + return err + } + if err := apply(ctx, q, locked); err != nil { + return err + } + return PruneChanges(ctx, q, session) + }) +} diff --git a/services/core/internal/persistence/postgres/sessionpg/session_test.go b/services/core/internal/persistence/postgres/sessionpg/session_test.go new file mode 100644 index 000000000..51d3e56e7 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/session_test.go @@ -0,0 +1,296 @@ +package sessionpg + +import ( + "context" + "errors" + "reflect" + "strings" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + 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/pgtest" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var errRollback = errors.New("roll back") + +func pgID(id uuid.UUID) pgtype.UUID { return pgtype.UUID{Bytes: id, Valid: true} } + +func exec(t *testing.T, pool *pgxpool.Pool, sql string, args ...any) { + t.Helper() + if _, err := pool.Exec(t.Context(), sql, args...); err != nil { + t.Fatal(err) + } +} + +// newEnvironment stores a fresh tenant's Session whose Environment has kind +// and status, and returns the tenant, Session and Environment IDs. +func newEnvironment(t *testing.T, pool *pgxpool.Pool, kind, status string) (pgtype.UUID, pgtype.UUID, pgtype.UUID) { + t.Helper() + tenant, session, environment := uuid.New(), uuid.New(), uuid.New() + exec(t, pool, `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash, configuration) + VALUES ($1, $2, 'codex', 'key', 'hash', jsonb_build_object('environment', jsonb_build_object('type', $3::text)))`, session, tenant, kind) + exec(t, pool, `INSERT INTO environments(id, session_id, status) VALUES ($1, $2, $3)`, environment, session, status) + return pgID(tenant), pgID(session), pgID(environment) +} + +func reserveInput(t *testing.T, pool *pgxpool.Pool, session pgtype.UUID) { + t.Helper() + exec(t, pool, `INSERT INTO environment_input_reservations(id, session_id, idempotency_key, batch, created_at, deadline) + VALUES ($1, $2, 'input', '[{}]', clock_timestamp(), clock_timestamp() + interval '1 hour')`, uuid.New(), session) +} + +func kinds(changes []sessions.SessionChange) []string { + var types []string + for _, change := range changes { + types = append(types, change.Event.Type) + } + return types +} + +func TestWithSessionLocksAppliesAndPrunes(t *testing.T) { + pool := pgtest.Open(t) + runner := pgunit.NewPool(pool) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "connected") + change := sessions.SessionChange{Event: v1.SessionEvent{Type: "agent.session.idle"}} + + err := WithSession(t.Context(), runner, pgID(uuid.New()), session, func(context.Context, *sqlc.Queries, sessions.LockedSession) error { + t.Fatal("applied to another tenant's Session") + return nil + }) + if !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("other tenant: %v", err) + } + + err = WithSession(t.Context(), runner, tenant, session, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { + if err := AppendChanges(ctx, q, session, change); err != nil { + return err + } + return errRollback + }) + if _, changes := journal(t, pool, session); !errors.Is(err, errRollback) || len(changes) != 0 { + t.Fatalf("failed apply committed %d changes: %v", len(changes), err) + } + + err = WithSession(t.Context(), runner, tenant, session, func(ctx context.Context, q *sqlc.Queries, locked sessions.LockedSession) error { + if locked.Deleted { + t.Fatal("live Session locked as deleted") + } + for range sessions.RetainedChanges + 2 { + if err := AppendChanges(ctx, q, session, change); err != nil { + return err + } + } + return nil + }) + if _, changes := journal(t, pool, session); err != nil || len(changes) != sessions.RetainedChanges { + t.Fatalf("retained %d changes: %v", len(changes), err) + } + + exec(t, pool, `UPDATE sessions SET deleted_at = clock_timestamp() WHERE id = $1`, session) + err = WithSession(t.Context(), runner, tenant, session, func(_ context.Context, _ *sqlc.Queries, locked sessions.LockedSession) error { + return locked.Public() + }) + if !errors.Is(err, sessions.ErrNotFound) { + t.Fatalf("deleted Session: %v", err) + } +} + +// WithSession holds the Session lock until its transaction ends: another +// transaction cannot lock the Session meanwhile, and can once it has ended. +func TestWithSessionHoldsTheSessionLock(t *testing.T) { + pool := pgtest.Open(t) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "connected") + lock := func(ctx context.Context) error { + return pgx.BeginFunc(ctx, pool, func(tx pgx.Tx) error { + if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '100ms'`); err != nil { + return err + } + _, err := LockSession(ctx, sqlc.New(tx), tenant, session) + return err + }) + } + + err := WithSession(t.Context(), pgunit.NewPool(pool), tenant, session, func(ctx context.Context, _ *sqlc.Queries, _ sessions.LockedSession) error { + return lock(ctx) + }) + var pgErr *pgconn.PgError + if !errors.As(err, &pgErr) || pgErr.Code != "55P03" { + t.Fatalf("locked a held Session: %v", err) + } + if err := lock(t.Context()); err != nil { + t.Fatalf("lock after the holder ended: %v", err) + } +} + +// A binding reads and writes only its tenant's Session's own Environment: +// another tenant, or another Session's Environment, is an error and changes no +// row. +func TestBindingScopesTheEnvironmentToItsSession(t *testing.T) { + pool := pgtest.Open(t) + tenant, session, environment := newEnvironment(t, pool, "openai_hosted", "connected") + _, _, foreign := newEnvironment(t, pool, "openai_hosted", "connected") + stranger := pgID(uuid.New()) + inTx(t, pool, func(q *sqlc.Queries) error { + if _, err := BindSession(q, stranger, session).LoadEnvironment(t.Context()); !errors.Is(err, sessions.ErrNotFound) { + t.Errorf("other tenant loaded the Environment: %v", err) + } + for name, test := range map[string]struct { + bound *SessionTx + environment pgtype.UUID + }{ + "other Session's Environment": {BindSession(q, tenant, session), foreign}, + "other tenant": {BindSession(q, stranger, session), environment}, + } { + id := uuid.UUID(test.environment.Bytes).String() + if _, err := test.bound.RecordEnvironmentFailure(t.Context(), id, "reason", nil); !errors.Is(err, pgx.ErrNoRows) { + t.Errorf("%s failed: %v", name, err) + } + if err := test.bound.ExpireEnvironment(t.Context(), id); !errors.Is(err, sessions.ErrNotFound) { + t.Errorf("%s expired: %v", name, err) + } + } + return nil + }) + for _, id := range []pgtype.UUID{environment, foreign} { + var status string + if err := pool.QueryRow(t.Context(), `SELECT status FROM environments WHERE id = $1`, id).Scan(&status); err != nil || status != "connected" { + t.Fatalf("Environment %s: %v", status, err) + } + } +} + +type environmentState struct { + environment, turn, input string + events int +} + +func readEnvironmentState(t *testing.T, pool *pgxpool.Pool, session pgtype.UUID) environmentState { + t.Helper() + var state environmentState + if err := pool.QueryRow(t.Context(), `SELECT e.status, COALESCE((SELECT status FROM turns WHERE session_id = e.session_id), ''), + (SELECT state FROM environment_input_reservations WHERE session_id = e.session_id), + (SELECT count(*) FROM session_events WHERE session_id = e.session_id) + FROM environments e WHERE e.session_id = $1`, session).Scan(&state.environment, &state.turn, &state.input, &state.events); err != nil { + t.Fatal(err) + } + return state +} + +// The binding writes only through its caller's transaction: rolling that +// transaction back reverts every participant's writes, and committing it keeps +// them with the events in the procedure's order. +func TestFailEnvironmentThroughTheBindingCommitsWithItsTransaction(t *testing.T) { + pool := pgtest.Open(t) + tenant, session, _ := newEnvironment(t, pool, "openai_hosted", "connected") + exec(t, pool, `INSERT INTO turns(id, session_id) VALUES ($1, $2)`, uuid.New(), session) + reserveInput(t, pool, session) + const reason = "Failed to provision environment" + fail := func(tx pgx.Tx) error { + bound := BindSession(sqlc.New(tx), tenant, session) + environment, err := bound.LoadEnvironment(t.Context()) + if err != nil { + return err + } + return sessions.FailEnvironment(t.Context(), bound, environment, reason, &sessions.ProvisioningFailureDetail{}) + } + + err := pgx.BeginFunc(t.Context(), pool, func(tx pgx.Tx) error { + if err := fail(tx); err != nil { + return err + } + return errRollback + }) + if got := readEnvironmentState(t, pool, session); !errors.Is(err, errRollback) || got != (environmentState{"connected", sessions.TurnQueued, sessions.EnvironmentInputPending, 0}) { + t.Fatalf("rolled back to %+v: %v", got, err) + } + + if err := pgx.BeginFunc(t.Context(), pool, fail); err != nil { + t.Fatal(err) + } + if got := readEnvironmentState(t, pool, session); got != (environmentState{"failed", sessions.TurnCancelled, sessions.EnvironmentInputFailed, 5}) { + t.Fatalf("committed %+v", got) + } + var failedAt time.Time + if err := pool.QueryRow(t.Context(), `SELECT failed_at FROM environments WHERE session_id = $1`, session).Scan(&failedAt); err != nil { + t.Fatal(err) + } + _, changes := journal(t, pool, session) + want := []string{"agent.session.environment.failed", "agent.session.turn.cancelled", "agent.session.idle", "error", "agent.session.failed"} + failed := changes[len(changes)-1] + if !reflect.DeepEqual(kinds(changes), want) || !failed.Settled || !failed.EnvironmentFailure.FailedAt.Equal(failedAt) || + failed.EnvironmentInputActivity == nil || failed.EnvironmentInputActivity.Status != "failed" { + t.Fatalf("journal %q, failed %+v", kinds(changes), failed) + } +} + +// Expired compute expires a live Environment, fails its pending input and +// reports the input activity that changed. +func TestTerminateEnvironmentThroughTheBindingExpires(t *testing.T) { + pool := pgtest.Open(t) + tenant, session, _ := newEnvironment(t, pool, "self_hosted", "connected") + reserveInput(t, pool, session) + inTx(t, pool, func(q *sqlc.Queries) error { + return sessions.TerminateEnvironment(t.Context(), BindSession(q, tenant, session), true, "reason", nil) + }) + if got := readEnvironmentState(t, pool, session); got != (environmentState{"expired", "", sessions.EnvironmentInputFailed, 1}) { + t.Fatalf("terminated %+v", got) + } + _, changes := journal(t, pool, session) + if changes[0].Event.Type != "agent.session.failed" || !changes[0].Settled || changes[0].EnvironmentInputActivity.Failure != "environment_unavailable" || changes[0].EnvironmentFailure != nil { + t.Fatalf("change %+v", changes[0]) + } +} + +func TestCreateEnvironmentDeviceBindsOneDevice(t *testing.T) { + pool := pgtest.Open(t) + hash := strings.Repeat("a", 64) + create := func(tenant, session, environment pgtype.UUID, device uuid.UUID, commit bool) error { + dedicated := sessions.ExecutionDevice{ID: device.String(), Name: "runtime", EnvironmentID: uuid.UUID(environment.Bytes).String()} + return pgx.BeginFunc(t.Context(), pool, func(tx pgx.Tx) error { + if err := sessions.CreateEnvironmentDevice(t.Context(), BindSession(sqlc.New(tx), tenant, session), dedicated, hash); err != nil || commit { + return err + } + return errRollback + }) + } + bound := func(session pgtype.UUID) []uuid.UUID { + rows, err := pool.Query(t.Context(), `SELECT device_id FROM session_devices WHERE session_id = $1`, session) + if err != nil { + t.Fatal(err) + } + devices, err := pgx.CollectRows(rows, pgx.RowTo[uuid.UUID]) + if err != nil { + t.Fatal(err) + } + return devices + } + + tenant, session, environment := newEnvironment(t, pool, "openai_hosted", "pending") + if err := create(tenant, session, environment, uuid.New(), false); !errors.Is(err, errRollback) || len(bound(session)) != 0 { + t.Fatalf("rolled back creation bound %v: %v", bound(session), err) + } + if err := create(pgID(uuid.New()), session, environment, uuid.New(), true); !errors.Is(err, sessions.ErrDeviceBindingConflict) || len(bound(session)) != 0 { + t.Fatalf("other tenant bound %v: %v", bound(session), err) + } + device := uuid.New() + if err := create(tenant, session, environment, device, true); err != nil || !reflect.DeepEqual(bound(session), []uuid.UUID{device}) { + t.Fatalf("bound %v: %v", bound(session), err) + } + if err := create(tenant, session, environment, uuid.New(), true); !errors.Is(err, sessions.ErrDeviceBindingConflict) { + t.Fatalf("second device: %v", err) + } + tenant, session, environment = newEnvironment(t, pool, "self_hosted", "pending") + if err := create(tenant, session, environment, uuid.New(), true); !errors.Is(err, sessions.ErrDeviceBindingConflict) { + t.Fatalf("self-hosted device: %v", err) + } +} diff --git a/services/core/internal/sessions/admission.go b/services/core/internal/sessions/admission.go new file mode 100644 index 000000000..279821331 --- /dev/null +++ b/services/core/internal/sessions/admission.go @@ -0,0 +1,71 @@ +package sessions + +import "context" + +// ComputeAdmissionTx reads the Session's managed compute for +// CheckComputeAdmission. +type ComputeAdmissionTx interface { + // LoadComputeSuspension reports whether the Session's managed compute is in + // its suspension cycle: quiescing, suspending, suspended, restoring or + // waking. + LoadComputeSuspension(ctx context.Context) (bool, error) +} + +// CheckComputeAdmission admits a new execution owner, such as a Turn that +// starts or a file write, only while the Session's managed compute is running +// or has no suspension; otherwise it is ErrTurnConflict. Existing receipts stay +// readable in every phase. +func CheckComputeAdmission(ctx context.Context, tx ComputeAdmissionTx) error { + suspended, err := tx.LoadComputeSuspension(ctx) + if err != nil { + return err + } + if suspended { + return ErrTurnConflict + } + return nil +} + +// FileWriteGateTx reads the Session's Environment file writes for +// CheckFileWriteGate. +type FileWriteGateTx interface { + // LoadPendingFileWrite reports whether a file write to the Session's + // Environment is pending. + LoadPendingFileWrite(ctx context.Context) (bool, error) +} + +// CheckFileWriteGate holds new work while a file write to the Session's +// Environment is pending: it is ErrTurnConflict. +func CheckFileWriteGate(ctx context.Context, tx FileWriteGateTx) error { + pending, err := tx.LoadPendingFileWrite(ctx) + if err != nil { + return err + } + if pending { + return ErrTurnConflict + } + return nil +} + +// InputStartTx is the Session transaction CheckInputStart runs in. +type InputStartTx interface { + FileWriteGateTx + ActiveTurnTx +} + +// CheckInputStart lets reserved Environment input, or a file write, start only +// on an idle Session: a pending file write, checked first, and an active Turn +// are both ErrTurnConflict. +func CheckInputStart(ctx context.Context, tx InputStartTx) error { + if err := CheckFileWriteGate(ctx, tx); err != nil { + return err + } + _, active, err := tx.LoadActiveTurn(ctx) + if err != nil { + return err + } + if active { + return ErrTurnConflict + } + return nil +} diff --git a/services/core/internal/sessions/admission_test.go b/services/core/internal/sessions/admission_test.go new file mode 100644 index 000000000..1d69f0bdc --- /dev/null +++ b/services/core/internal/sessions/admission_test.go @@ -0,0 +1,43 @@ +package sessions + +import ( + "errors" + "testing" +) + +func TestCheckComputeAdmission(t *testing.T) { + for name, test := range map[string]struct { + load func() (bool, error) + want error + }{ + "running": {returns(false), nil}, + "suspended": {returns(true), ErrTurnConflict}, + "failing": {func() (bool, error) { return false, errStorage }, errStorage}, + } { + f := &fakeTx{t: t, loadComputeSuspension: test.load} + if err := CheckComputeAdmission(t.Context(), f); !errors.Is(err, test.want) { + t.Errorf("%s: %v", name, err) + } + assertCalls(t, f, "LoadComputeSuspension") + } +} + +// A pending file write is checked before the active Turn. +func TestCheckInputStart(t *testing.T) { + running := turnWith(TurnInProgress) + for name, test := range map[string]struct { + tx *fakeTx + want error + calls []string + }{ + "idle": {&fakeTx{t: t, loadPendingFileWrite: returns(false), loadActiveTurn: activeTurn(nil)}, nil, []string{"LoadPendingFileWrite", "LoadActiveTurn"}}, + "file write": {&fakeTx{t: t, loadPendingFileWrite: returns(true)}, ErrTurnConflict, []string{"LoadPendingFileWrite"}}, + "active Turn": {&fakeTx{t: t, loadPendingFileWrite: returns(false), loadActiveTurn: activeTurn(&running)}, ErrTurnConflict, []string{"LoadPendingFileWrite", "LoadActiveTurn"}}, + "failing write": {&fakeTx{t: t, loadPendingFileWrite: func() (bool, error) { return false, errStorage }}, errStorage, []string{"LoadPendingFileWrite"}}, + } { + if err := CheckInputStart(t.Context(), test.tx); !errors.Is(err, test.want) { + t.Errorf("%s: %v", name, err) + } + assertCalls(t, test.tx, test.calls...) + } +} diff --git a/services/core/internal/sessions/cancellation.go b/services/core/internal/sessions/cancellation.go new file mode 100644 index 000000000..7523e61ae --- /dev/null +++ b/services/core/internal/sessions/cancellation.go @@ -0,0 +1,105 @@ +package sessions + +import "context" + +// turnCancellation is what a cancellation request settles beyond recording the +// request, decided from the Turn as it was before the request. +type turnCancellation int + +const ( + // cancellationRequested only records the request: a running Turn ends when + // its execution stops, and a repeated request adds nothing. + cancellationRequested turnCancellation = iota + // cancellationEndsTurn ends a queued Turn, which never started, as + // cancelled. + cancellationEndsTurn + // cancellationReportsActivity reports the Session activity of a waiting + // Turn on its first request: it no longer waits on actions. + cancellationReportsActivity +) + +// decideCancellation decides what requesting the cancellation of turn, read +// before the request, settles. +func decideCancellation(turn Turn) turnCancellation { + switch { + case turn.Status == TurnQueued: + return cancellationEndsTurn + case turn.Status == TurnWaiting && turn.CancelRequestedAt.IsZero(): + return cancellationReportsActivity + } + return cancellationRequested +} + +// TurnCancellationTx is the Session transaction CancelTurn runs in. +type TurnCancellationTx interface { + JournalTx + // RequestTurnCancel records a cancellation request on an active Turn of + // the Session. A queued Turn becomes cancelled; the database clock stamps + // the request and the end. + RequestTurnCancel(ctx context.Context, turn string) error + // LoadTurn reads one of the Session's Turns. + LoadTurn(ctx context.Context, turn string) (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 +} + +// CancelTurn requests the cancellation of turn, the Session's active Turn read +// under the Session lock, and publishes what the request settles: a queued Turn +// ends cancelled, and a waiting Turn reports its Session activity. Both read +// the Turn again after the request, which set its timestamps. +func CancelTurn(ctx context.Context, tx TurnCancellationTx, turn Turn) error { + cancellation := decideCancellation(turn) + if err := tx.RequestTurnCancel(ctx, turn.ID); err != nil { + return err + } + switch cancellation { + case cancellationEndsTurn: + cancelled, err := tx.LoadTurn(ctx, turn.ID) + if err != nil { + return err + } + ending, err := tx.LoadEnding(ctx, turn.ID) + if err != nil { + return err + } + return tx.ApplyTurnEnd(ctx, turn.ID, EndTurn(cancelled, ending)) + case cancellationReportsActivity: + cancelling, err := tx.LoadTurn(ctx, turn.ID) + if err != nil { + return err + } + usage, err := tx.LoadUsage(ctx) + if err != nil { + return err + } + return tx.AppendChanges(ctx, ActivityChange(cancelling, usage, nil)) + } + return nil +} + +// CancellationTx is the Session transaction CancelWork runs in. +type CancellationTx interface { + TurnCancellationTx + ActiveTurnTx + // CancelPendingInput settles the Session's pending Environment input + // reservation as cancelled. + CancelPendingInput(ctx context.Context) error +} + +// CancelWork cancels the Session's active work: it requests the cancellation +// of the active Turn, then cancels the pending Environment input. Cancellation +// requests do not prove that native work has stopped. +func CancelWork(ctx context.Context, tx CancellationTx) error { + turn, active, err := tx.LoadActiveTurn(ctx) + if err != nil { + return err + } + if active { + if err := CancelTurn(ctx, tx, turn); err != nil { + return err + } + } + return tx.CancelPendingInput(ctx) +} diff --git a/services/core/internal/sessions/cancellation_test.go b/services/core/internal/sessions/cancellation_test.go new file mode 100644 index 000000000..624b83202 --- /dev/null +++ b/services/core/internal/sessions/cancellation_test.go @@ -0,0 +1,87 @@ +package sessions + +import ( + "encoding/json" + "errors" + "reflect" + "testing" + "time" +) + +func TestDecideCancellation(t *testing.T) { + requested := turnWith(TurnWaiting) + requested.CancelRequestedAt = time.Unix(1700000000, 0) + for name, test := range map[string]struct { + turn Turn + want turnCancellation + }{ + "queued": {turnWith(TurnQueued), cancellationEndsTurn}, + "waiting": {turnWith(TurnWaiting), cancellationReportsActivity}, + "waiting requested": {requested, cancellationRequested}, + "in progress": {turnWith(TurnInProgress), cancellationRequested}, + } { + if got := decideCancellation(test.turn); got != test.want { + t.Errorf("%s: %v", name, got) + } + } +} + +// CancelTurn decides from the Turn read before the request, and settles with +// the Turn read after it. +func TestCancelTurn(t *testing.T) { + usage := json.RawMessage(`{"input_tokens":1}`) + t.Run("queued ends cancelled", func(t *testing.T) { + cancelled, ending := turnWith(TurnCancelled), Ending{Usage: usage} + var ends []TurnEnd + f := &fakeTx{t: t, requestTurnCancel: done, loadTurn: returns(cancelled), loadEnding: returns(ending), + applyTurnEnd: func(end TurnEnd) error { ends = append(ends, end); return nil }} + if err := CancelTurn(t.Context(), f, turnWith(TurnQueued)); err != nil { + t.Fatal(err) + } + assertCalls(t, f, "RequestTurnCancel "+testTurn, "LoadTurn "+testTurn, "LoadEnding "+testTurn, "ApplyTurnEnd "+testTurn) + if !reflect.DeepEqual(ends, []TurnEnd{EndTurn(cancelled, ending)}) { + t.Fatalf("end %+v", ends) + } + }) + t.Run("waiting reports activity", func(t *testing.T) { + cancelling := turnWith(TurnWaiting) + cancelling.CancelRequestedAt = time.Unix(1700000001, 0) + var changes []SessionChange + f := &fakeTx{t: t, requestTurnCancel: done, loadTurn: returns(cancelling), loadUsage: returns(usage), appendChanges: collect(&changes)} + if err := CancelTurn(t.Context(), f, turnWith(TurnWaiting)); err != nil { + t.Fatal(err) + } + assertCalls(t, f, "RequestTurnCancel "+testTurn, "LoadTurn "+testTurn, "LoadUsage", "AppendChanges agent.session.in_progress") + if !reflect.DeepEqual(changes, []SessionChange{ActivityChange(cancelling, usage, nil)}) { + t.Fatalf("changes %+v", changes) + } + }) + t.Run("in progress only requests", func(t *testing.T) { + f := &fakeTx{t: t, requestTurnCancel: done} + if err := CancelTurn(t.Context(), f, turnWith(TurnInProgress)); err != nil { + t.Fatal(err) + } + assertCalls(t, f, "RequestTurnCancel "+testTurn) + }) +} + +func TestCancelWork(t *testing.T) { + idle := &fakeTx{t: t, loadActiveTurn: activeTurn(nil), cancelPendingInput: done} + if err := CancelWork(t.Context(), idle); err != nil { + t.Fatal(err) + } + assertCalls(t, idle, "LoadActiveTurn", "CancelPendingInput") + + running := turnWith(TurnInProgress) + active := &fakeTx{t: t, loadActiveTurn: activeTurn(&running), requestTurnCancel: done, cancelPendingInput: done} + if err := CancelWork(t.Context(), active); err != nil { + t.Fatal(err) + } + assertCalls(t, active, "LoadActiveTurn", "RequestTurnCancel "+testTurn, "CancelPendingInput") + + failing := &fakeTx{t: t, loadActiveTurn: activeTurn(&running), requestTurnCancel: func() error { return errStorage }} + if err := CancelWork(t.Context(), failing); !errors.Is(err, errStorage) { + t.Fatal(err) + } + assertCalls(t, failing, "LoadActiveTurn", "RequestTurnCancel "+testTurn) +} diff --git a/services/core/internal/sessions/doc.go b/services/core/internal/sessions/doc.go index d40126ef8..b3f93cc48 100644 --- a/services/core/internal/sessions/doc.go +++ b/services/core/internal/sessions/doc.go @@ -4,4 +4,8 @@ // operations return. It also decides what Session writes publish: the public // changes that report Turn and Session transitions, what a Turn that ends // settles, measured Turn usage and the Session activity each change reports. +// 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 and the admission +// gates. package sessions diff --git a/services/core/internal/sessions/environment_device.go b/services/core/internal/sessions/environment_device.go new file mode 100644 index 000000000..ee6377390 --- /dev/null +++ b/services/core/internal/sessions/environment_device.go @@ -0,0 +1,31 @@ +package sessions + +import "context" + +// EnvironmentDeviceTx is the Session transaction CreateEnvironmentDevice runs +// in. +type EnvironmentDeviceTx interface { + // LoadBoundDevice reports whether the Session is bound to a Runtime device + // that still has authority. + LoadBoundDevice(ctx context.Context) (bool, error) + // InsertEnvironmentDevice inserts device, with the credential hash, as the + // dedicated Runtime device of the Session's hosted Environment and binds + // it to the Session. An Environment that already has a device or cannot + // take one is ErrDeviceBindingConflict. + InsertEnvironmentDevice(ctx context.Context, device ExecutionDevice, credentialHash string) error +} + +// CreateEnvironmentDevice creates the dedicated Runtime device of the +// Session's hosted Environment and binds it to the Session. A Session that is +// already bound to a device is ErrDeviceBindingConflict, so an existing +// credential is never widened. +func CreateEnvironmentDevice(ctx context.Context, tx EnvironmentDeviceTx, device ExecutionDevice, credentialHash string) error { + bound, err := tx.LoadBoundDevice(ctx) + if err != nil { + return err + } + if bound { + return ErrDeviceBindingConflict + } + return tx.InsertEnvironmentDevice(ctx, device, credentialHash) +} diff --git a/services/core/internal/sessions/environment_device_test.go b/services/core/internal/sessions/environment_device_test.go new file mode 100644 index 000000000..14e9e4088 --- /dev/null +++ b/services/core/internal/sessions/environment_device_test.go @@ -0,0 +1,22 @@ +package sessions + +import ( + "errors" + "testing" +) + +// A Session that is already bound to a device never gets a second one. +func TestCreateEnvironmentDevice(t *testing.T) { + device := ExecutionDevice{ID: "device", Name: "runtime", EnvironmentID: "environment"} + free := &fakeTx{t: t, loadBoundDevice: returns(false), insertEnvironmentDevice: done} + if err := CreateEnvironmentDevice(t.Context(), free, device, "hash"); err != nil { + t.Fatal(err) + } + assertCalls(t, free, "LoadBoundDevice", "InsertEnvironmentDevice device runtime environment hash") + + bound := &fakeTx{t: t, loadBoundDevice: returns(true)} + if err := CreateEnvironmentDevice(t.Context(), bound, device, "hash"); !errors.Is(err, ErrDeviceBindingConflict) { + t.Fatal(err) + } + assertCalls(t, bound, "LoadBoundDevice") +} diff --git a/services/core/internal/sessions/environment_termination.go b/services/core/internal/sessions/environment_termination.go new file mode 100644 index 000000000..2707b7457 --- /dev/null +++ b/services/core/internal/sessions/environment_termination.go @@ -0,0 +1,165 @@ +package sessions + +import ( + "context" + "encoding/json" + "errors" + "time" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" +) + +// EnvironmentType returns the type an Environment configuration snapshot +// names: self_hosted or openai_hosted. +func EnvironmentType(configuration json.RawMessage) (string, error) { + var config struct { + Type string `json:"type"` + } + if err := json.Unmarshal(configuration, &config); err != nil || (config.Type != "self_hosted" && config.Type != "openai_hosted") { + return "", errors.New("invalid stored Environment type") + } + return config.Type, nil +} + +// EnvironmentStateChange is the pinned agent.session.environment. +// event for an Environment status the transaction has written. A failure +// carries the observed official connection error; the failed step travels only +// in the separate error event and the Session's error. +func EnvironmentStateChange(environment, kind, status string) SessionChange { + state := &v1.SessionEnvironmentState{ID: environment, Type: kind, Status: status} + if status == "failed" { + state.Error = &v1.StreamError{Type: "environment_error", Code: "environment_connection_failed", Message: "The environment failed to connect."} + } + return SessionChange{Event: v1.SessionEvent{Type: "agent.session.environment." + status, Environment: state}} +} + +// environmentFailedChanges end a Session whose Environment failed: an error +// event carrying the safe reason, then one agent.session.failed snapshot of +// the input activity, the Session's usage and the failure, matching later +// Session reads. +func environmentFailedChanges(failure EnvironmentFailure, activity *EnvironmentInputActivity, pending bool, usage json.RawMessage) []SessionChange { + return []SessionChange{ + {Event: v1.SessionEvent{Type: "error", Error: &v1.StreamError{Type: "environment_error", Code: "sandbox_error", Message: failure.Reason}}}, + { + Event: v1.SessionEvent{Type: "agent.session.failed"}, + EnvironmentInputActivity: activity, EnvironmentFailure: &EnvironmentFailure{Reason: failure.Reason, FailedAt: failure.FailedAt}, + SessionUsage: usage, Settled: !pending, + }, + } +} + +// EnvironmentFailureTx is the Session transaction FailEnvironment runs in. +type EnvironmentFailureTx interface { + CancellationTx + InputActivityTx + // RecordEnvironmentFailure marks the Session's live Environment failed + // with the reason and the private detail, and returns the failure time the + // database clock recorded. + RecordEnvironmentFailure(ctx context.Context, environment, reason string, detail *ProvisioningFailureDetail) (time.Time, error) + // FailPendingInput settles the Session's pending Environment input + // reservation as failed. + FailPendingInput(ctx context.Context) error +} + +// FailEnvironment records the failure of the Session's live environment and +// ends the Session failed, in the observed official order: +// agent.session.environment.failed, the changes of cancelling the Session's +// work, an error event with the safe reason, then one agent.session.failed +// snapshot. Pending input settles as failed first, so the snapshot reports the +// settled input activity. +func FailEnvironment(ctx context.Context, tx EnvironmentFailureTx, environment Environment, reason string, detail *ProvisioningFailureDetail) error { + failedAt, err := tx.RecordEnvironmentFailure(ctx, environment.ID, reason, detail) + if err != nil { + return err + } + kind, err := EnvironmentType(environment.Configuration) + if err != nil { + return err + } + if err := tx.AppendChanges(ctx, EnvironmentStateChange(environment.ID, kind, "failed")); err != nil { + return err + } + if err := tx.FailPendingInput(ctx); err != nil { + return err + } + if err := CancelWork(ctx, tx); err != nil { + return err + } + state, err := tx.LoadEnvironmentInput(ctx) + if err != nil { + return err + } + activity, pending := InputActivity(state) + usage, err := tx.LoadUsage(ctx) + if err != nil { + return err + } + return tx.AppendChanges(ctx, environmentFailedChanges(EnvironmentFailure{Reason: reason, FailedAt: failedAt}, activity, pending, usage)...) +} + +// environmentTermination is how the end of a Session's managed compute settles +// its Environment. +type environmentTermination int + +const ( + // terminationFails records the provisioning failure of a live Environment + // whose compute did not expire. + terminationFails environmentTermination = iota + // terminationExpires expires a live Environment whose compute expired and + // settles its input. Public expiry neither asserts compute removal nor + // invents an expired event. + terminationExpires + // terminationSettles only settles the remaining input of an Environment + // that already failed or expired, without repeating its events. + terminationSettles +) + +// decideEnvironmentTermination decides how the end of managed compute settles +// an Environment of status whose compute did or did not expire. +func decideEnvironmentTermination(status string, computeExpired bool) environmentTermination { + switch { + case status == "failed" || status == "expired": + return terminationSettles + case computeExpired: + return terminationExpires + } + return terminationFails +} + +// EnvironmentTerminationTx is the Session transaction TerminateEnvironment +// runs in. +type EnvironmentTerminationTx interface { + EnvironmentFailureTx + // LoadEnvironment reads the Session's Environment. + LoadEnvironment(ctx context.Context) (Environment, error) + // ExpireEnvironment sets the status of the Session's Environment to + // expired. + ExpireEnvironment(ctx context.Context, environment string) error +} + +// TerminateEnvironment settles the Session's Environment when its managed +// compute ends. A live Environment whose compute did not expire fails with the +// reason and detail, as FailEnvironment reports it. Otherwise the Environment +// expires if it is still live, its pending input fails and its work is +// cancelled, and the input activity this changes is reported. +func TerminateEnvironment(ctx context.Context, tx EnvironmentTerminationTx, computeExpired bool, reason string, detail *ProvisioningFailureDetail) error { + environment, err := tx.LoadEnvironment(ctx) + if err != nil { + return err + } + termination := decideEnvironmentTermination(environment.Status, computeExpired) + if termination == terminationFails { + return FailEnvironment(ctx, tx, environment, reason, detail) + } + return TrackInputActivity(ctx, tx, func(ctx context.Context) error { + if termination == terminationExpires { + if err := tx.ExpireEnvironment(ctx, environment.ID); err != nil { + return err + } + } + if err := tx.FailPendingInput(ctx); err != nil { + return err + } + return CancelWork(ctx, tx) + }) +} diff --git a/services/core/internal/sessions/environment_termination_test.go b/services/core/internal/sessions/environment_termination_test.go new file mode 100644 index 000000000..4f86127c2 --- /dev/null +++ b/services/core/internal/sessions/environment_termination_test.go @@ -0,0 +1,133 @@ +package sessions + +import ( + "encoding/json" + "reflect" + "testing" + "time" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" +) + +func TestEnvironmentType(t *testing.T) { + for configuration, want := range map[string]string{ + `{"type":"self_hosted"}`: "self_hosted", + `{"type":"openai_hosted"}`: "openai_hosted", + `{"type":"cloud"}`: "", + `{}`: "", + `[`: "", + } { + got, err := EnvironmentType(json.RawMessage(configuration)) + if got != want || (err == nil) != (want != "") { + t.Errorf("%s: %q, %v", configuration, got, err) + } + } +} + +func TestEnvironmentStateChange(t *testing.T) { + connected := EnvironmentStateChange("environment", "self_hosted", "connected") + if connected.Event.Type != "agent.session.environment.connected" || !reflect.DeepEqual(connected.Event.Environment, &v1.SessionEnvironmentState{ID: "environment", Type: "self_hosted", Status: "connected"}) { + t.Fatalf("connected %+v", connected.Event) + } + failed := EnvironmentStateChange("environment", "openai_hosted", "failed") + if failed.Event.Type != "agent.session.environment.failed" || failed.Event.Environment.Error == nil || failed.Event.Environment.Error.Code != "environment_connection_failed" { + t.Fatalf("failed %+v", failed.Event) + } +} + +func TestDecideEnvironmentTermination(t *testing.T) { + for _, test := range []struct { + status string + expired bool + want environmentTermination + }{ + {"connected", false, terminationFails}, + {"disconnected", true, terminationExpires}, + {"failed", false, terminationSettles}, + {"failed", true, terminationSettles}, + {"expired", true, terminationSettles}, + } { + if got := decideEnvironmentTermination(test.status, test.expired); got != test.want { + t.Errorf("%s expired=%v: %v", test.status, test.expired, got) + } + } +} + +// FailEnvironment journals the Environment failure, then cancels the Session's +// work, then reports the error and the failed Session with the recorded +// failure time and the settled input activity. +func TestFailEnvironment(t *testing.T) { + running := turnWith(TurnInProgress) + failedAt := time.Unix(1700000000, 0) + settled := &EnvironmentInputState{State: EnvironmentInputFailed, EnvironmentType: "openai_hosted", SettledAt: failedAt} + usage := json.RawMessage(`{"input_tokens":1}`) + environment := Environment{ID: "environment", Configuration: json.RawMessage(`{"type":"openai_hosted"}`)} + var changes []SessionChange + f := &fakeTx{t: t, + recordEnvironmentFailure: returns(failedAt), appendChanges: collect(&changes), failPendingInput: done, + loadActiveTurn: activeTurn(&running), requestTurnCancel: done, cancelPendingInput: done, + loadEnvironmentInput: inputs(settled), loadUsage: returns(usage), + } + if err := FailEnvironment(t.Context(), f, environment, "Failed to provision environment", nil); err != nil { + t.Fatal(err) + } + assertCalls(t, f, + "RecordEnvironmentFailure environment Failed to provision environment", + "AppendChanges agent.session.environment.failed", + "FailPendingInput", + "LoadActiveTurn", "RequestTurnCancel "+testTurn, "CancelPendingInput", + "LoadEnvironmentInput", "LoadUsage", + "AppendChanges error,agent.session.failed", + ) + activity, _ := InputActivity(settled) + want := append([]SessionChange{EnvironmentStateChange("environment", "openai_hosted", "failed")}, + environmentFailedChanges(EnvironmentFailure{Reason: "Failed to provision environment", FailedAt: failedAt}, activity, false, usage)...) + if !reflect.DeepEqual(changes, want) || changes[1].Event.Error.Message != "Failed to provision environment" || !changes[2].Settled || changes[2].EnvironmentFailure.FailedAt != failedAt { + t.Fatalf("changes %+v", changes) + } + + invalid := &fakeTx{t: t, recordEnvironmentFailure: returns(failedAt)} + if err := FailEnvironment(t.Context(), invalid, Environment{ID: "environment", Configuration: json.RawMessage(`{}`)}, "reason", nil); err == nil { + t.Fatal("invalid Environment type failed") + } + assertCalls(t, invalid, "RecordEnvironmentFailure environment reason") +} + +func TestTerminateEnvironment(t *testing.T) { + configuration := json.RawMessage(`{"type":"openai_hosted"}`) + connected := Environment{ID: "environment", Status: "connected", Configuration: configuration} + var changes []SessionChange + live := &fakeTx{t: t, + loadEnvironment: returns(connected), recordEnvironmentFailure: returns(time.Unix(1700000000, 0)), appendChanges: collect(&changes), + failPendingInput: done, loadActiveTurn: activeTurn(nil), cancelPendingInput: done, + loadEnvironmentInput: inputs(), loadUsage: returns(json.RawMessage(`{}`)), + } + if err := TerminateEnvironment(t.Context(), live, false, "reason", nil); err != nil { + t.Fatal(err) + } + assertCalls(t, live, + "LoadEnvironment", "RecordEnvironmentFailure environment reason", + "AppendChanges agent.session.environment.failed", + "FailPendingInput", "LoadActiveTurn", "CancelPendingInput", + "LoadEnvironmentInput", "LoadUsage", + "AppendChanges error,agent.session.failed", + ) + + expired := &fakeTx{t: t, + loadEnvironment: returns(connected), loadEnvironmentInput: inputs(), expireEnvironment: done, + failPendingInput: done, loadActiveTurn: activeTurn(nil), cancelPendingInput: done, + } + if err := TerminateEnvironment(t.Context(), expired, true, "reason", nil); err != nil { + t.Fatal(err) + } + assertCalls(t, expired, "LoadEnvironment", "LoadEnvironmentInput", "ExpireEnvironment environment", "FailPendingInput", "LoadActiveTurn", "CancelPendingInput", "LoadEnvironmentInput") + + terminal := &fakeTx{t: t, + loadEnvironment: returns(Environment{ID: "environment", Status: "failed", Configuration: configuration}), + loadEnvironmentInput: inputs(), failPendingInput: done, loadActiveTurn: activeTurn(nil), cancelPendingInput: done, + } + if err := TerminateEnvironment(t.Context(), terminal, false, "reason", nil); err != nil { + t.Fatal(err) + } + assertCalls(t, terminal, "LoadEnvironment", "LoadEnvironmentInput", "FailPendingInput", "LoadActiveTurn", "CancelPendingInput", "LoadEnvironmentInput") +} diff --git a/services/core/internal/sessions/input_activity.go b/services/core/internal/sessions/input_activity.go new file mode 100644 index 000000000..bbaed3705 --- /dev/null +++ b/services/core/internal/sessions/input_activity.go @@ -0,0 +1,125 @@ +package sessions + +import ( + "context" + "encoding/json" + "time" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" +) + +// EnvironmentInputState is the Session's latest Environment input reservation +// that no Turn has admitted or superseded, with the facts of its Environment. +// While a Turn is active or newer than the latest reservation, the Session has +// none. +type EnvironmentInputState struct { + // State is the reservation's state. + State string + Initial bool + // CreatedAt is when the reservation was made, and SettledAt when it stopped + // being pending; SettledAt is zero while it is pending. + CreatedAt time.Time + SettledAt time.Time + // FailureCode is the recorded failure of a failed reservation, empty when + // none was recorded. + FailureCode string + EnvironmentID string + // EnvironmentType is the type the Session's Environment configuration names. + EnvironmentType string + // EnvironmentStatus is the Environment's connection status. + EnvironmentStatus string +} + +// InputActivity projects the Session's input activity from its latest +// reservation, nil when it has none, and reports whether that reservation is +// still pending and can start a Turn. A hosted initial input that is pending or +// cancelled has no activity: no Turn has started, and the pinned Session +// contract permits idle while a hosted Environment provisions, so neither a +// caller action nor an invented transition is reported. It still reports +// whether the input is pending. +func InputActivity(state *EnvironmentInputState) (*EnvironmentInputActivity, bool) { + if state == nil { + return nil, false + } + pending := state.State == EnvironmentInputPending + activity := &EnvironmentInputActivity{Status: "idle", LastActiveAt: state.CreatedAt} + if !state.SettledAt.IsZero() { + activity.LastActiveAt = state.SettledAt + } + if state.State == EnvironmentInputFailed { + activity.Status, activity.Failure = "failed", "environment_unavailable" + if state.FailureCode != "" { + activity.Failure = state.FailureCode + } + } + if state.Initial && state.State == EnvironmentInputExpired { + activity.Status = "failed" + } + if state.EnvironmentType == "openai_hosted" && state.Initial && + (state.State == EnvironmentInputPending || state.State == EnvironmentInputCancelled) { + return nil, pending + } + if pending && state.EnvironmentType != "openai_hosted" && state.EnvironmentStatus != "connected" { + activity.Status = "requires_action" + activity.EnvironmentID = state.EnvironmentID + } + return activity, pending +} + +// inputActivityChanged reports whether after is input activity the Session has +// not reported yet: it exists and differs from before in status, Environment or +// failure. Admitted input has no activity; its Turn and Session events report +// it. +func inputActivityChanged(before, after *EnvironmentInputActivity) bool { + if after == nil { + return false + } + return before == nil || before.Status != after.Status || before.EnvironmentID != after.EnvironmentID || before.Failure != after.Failure +} + +// inputActivityChange is the agent.session. snapshot that reports input +// activity with the Session's usage. It is settled once the reservation stops +// being pending. +func inputActivityChange(activity *EnvironmentInputActivity, pending bool, usage json.RawMessage) SessionChange { + return SessionChange{ + Event: v1.SessionEvent{Type: "agent.session." + activity.Status}, + EnvironmentInputActivity: activity, SessionUsage: usage, Settled: !pending, + } +} + +// InputActivityTx is the Session transaction TrackInputActivity runs in. +type InputActivityTx interface { + JournalTx + // 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) +} + +// TrackInputActivity runs apply, the caller's writes in the same transaction, +// and reports the input activity they change. It reads the activity before and +// after apply, and only when it changed reads the Session's usage and journals +// the snapshot. When apply fails, it returns at once and the caller's +// transaction rolls back. +func TrackInputActivity(ctx context.Context, tx InputActivityTx, apply func(context.Context) error) error { + state, err := tx.LoadEnvironmentInput(ctx) + if err != nil { + return err + } + before, _ := InputActivity(state) + if err := apply(ctx); err != nil { + return err + } + if state, err = tx.LoadEnvironmentInput(ctx); err != nil { + return err + } + after, pending := InputActivity(state) + if !inputActivityChanged(before, after) { + return nil + } + usage, err := tx.LoadUsage(ctx) + if err != nil { + return err + } + return tx.AppendChanges(ctx, inputActivityChange(after, pending, usage)) +} diff --git a/services/core/internal/sessions/input_activity_test.go b/services/core/internal/sessions/input_activity_test.go new file mode 100644 index 000000000..6855e6d38 --- /dev/null +++ b/services/core/internal/sessions/input_activity_test.go @@ -0,0 +1,103 @@ +package sessions + +import ( + "context" + "encoding/json" + "errors" + "reflect" + "testing" + "time" +) + +func TestInputActivity(t *testing.T) { + created, settled := time.Unix(1700000000, 0), time.Unix(1700000060, 0) + state := func(state, kind, status string, initial bool) *EnvironmentInputState { + s := &EnvironmentInputState{State: state, Initial: initial, CreatedAt: created, EnvironmentID: "environment", EnvironmentType: kind, EnvironmentStatus: status} + if state != EnvironmentInputPending { + s.SettledAt = settled + } + return s + } + coded := state(EnvironmentInputFailed, "self_hosted", "connected", false) + coded.FailureCode = "model_provider_required" + for name, test := range map[string]struct { + state *EnvironmentInputState + activity *EnvironmentInputActivity + pending bool + }{ + "none": {nil, nil, false}, + "pending disconnected": {state(EnvironmentInputPending, "self_hosted", "disconnected", false), &EnvironmentInputActivity{Status: "requires_action", EnvironmentID: "environment", LastActiveAt: created}, true}, + "pending connected": {state(EnvironmentInputPending, "self_hosted", "connected", false), &EnvironmentInputActivity{Status: "idle", LastActiveAt: created}, true}, + "pending hosted": {state(EnvironmentInputPending, "openai_hosted", "disconnected", false), &EnvironmentInputActivity{Status: "idle", LastActiveAt: created}, true}, + "failed": {state(EnvironmentInputFailed, "self_hosted", "connected", false), &EnvironmentInputActivity{Status: "failed", Failure: "environment_unavailable", LastActiveAt: settled}, false}, + "failed with code": {coded, &EnvironmentInputActivity{Status: "failed", Failure: "model_provider_required", LastActiveAt: settled}, false}, + "initial expired": {state(EnvironmentInputExpired, "self_hosted", "disconnected", true), &EnvironmentInputActivity{Status: "failed", LastActiveAt: settled}, false}, + "expired": {state(EnvironmentInputExpired, "self_hosted", "disconnected", false), &EnvironmentInputActivity{Status: "idle", LastActiveAt: settled}, false}, + "hosted initial pending": {state(EnvironmentInputPending, "openai_hosted", "disconnected", true), nil, true}, + "hosted initial cancelled": {state(EnvironmentInputCancelled, "openai_hosted", "disconnected", true), nil, false}, + "hosted initial failed reports": {state(EnvironmentInputFailed, "openai_hosted", "disconnected", true), &EnvironmentInputActivity{Status: "failed", Failure: "environment_unavailable", LastActiveAt: settled}, false}, + } { + activity, pending := InputActivity(test.state) + if !reflect.DeepEqual(activity, test.activity) || pending != test.pending { + t.Errorf("%s: %+v, %v", name, activity, pending) + } + } +} + +func TestInputActivityChanged(t *testing.T) { + idle := &EnvironmentInputActivity{Status: "idle"} + for name, test := range map[string]struct { + before, after *EnvironmentInputActivity + want bool + }{ + "admitted": {idle, nil, false}, + "new": {nil, idle, true}, + "same": {idle, &EnvironmentInputActivity{Status: "idle", LastActiveAt: time.Unix(1, 0)}, false}, + "status": {idle, &EnvironmentInputActivity{Status: "failed"}, true}, + "environment": {&EnvironmentInputActivity{Status: "requires_action", EnvironmentID: "a"}, &EnvironmentInputActivity{Status: "requires_action", EnvironmentID: "b"}, true}, + "failure": {&EnvironmentInputActivity{Status: "failed", Failure: "a"}, &EnvironmentInputActivity{Status: "failed", Failure: "b"}, true}, + } { + if got := inputActivityChanged(test.before, test.after); got != test.want { + t.Errorf("%s: %v", name, got) + } + } +} + +// TrackInputActivity reads the Session's usage and journals a snapshot only +// when apply changed the input activity. +func TestTrackInputActivity(t *testing.T) { + pending := &EnvironmentInputState{State: EnvironmentInputPending, EnvironmentID: "environment", EnvironmentType: "self_hosted", EnvironmentStatus: "disconnected"} + failed := &EnvironmentInputState{State: EnvironmentInputFailed, EnvironmentType: "self_hosted", SettledAt: time.Unix(1700000000, 0)} + usage := json.RawMessage(`{"input_tokens":1}`) + apply := func(f *fakeTx, err error) func(context.Context) error { + return func(context.Context) error { + f.calls = append(f.calls, "apply") + return err + } + } + + var changes []SessionChange + changed := &fakeTx{t: t, loadEnvironmentInput: inputs(pending, failed), loadUsage: returns(usage), appendChanges: collect(&changes)} + if err := TrackInputActivity(t.Context(), changed, apply(changed, nil)); err != nil { + t.Fatal(err) + } + assertCalls(t, changed, "LoadEnvironmentInput", "apply", "LoadEnvironmentInput", "LoadUsage", "AppendChanges agent.session.failed") + after, _ := InputActivity(failed) + if !reflect.DeepEqual(changes, []SessionChange{inputActivityChange(after, false, usage)}) || !changes[0].Settled { + t.Fatalf("changes %+v", changes) + } + + for name, states := range map[string][]*EnvironmentInputState{"unchanged": {pending, pending}, "admitted": {pending, nil}} { + f := &fakeTx{t: t, loadEnvironmentInput: inputs(states...)} + if err := TrackInputActivity(t.Context(), f, apply(f, nil)); err != nil { + t.Fatal(name, err) + } + assertCalls(t, f, "LoadEnvironmentInput", "apply", "LoadEnvironmentInput") + } + + failing := &fakeTx{t: t, loadEnvironmentInput: inputs(pending, failed)} + if err := TrackInputActivity(t.Context(), failing, apply(failing, errStorage)); !errors.Is(err, errStorage) { + t.Fatal(err) + } + assertCalls(t, failing, "LoadEnvironmentInput", "apply") +} diff --git a/services/core/internal/sessions/transaction.go b/services/core/internal/sessions/transaction.go new file mode 100644 index 000000000..10e58f34c --- /dev/null +++ b/services/core/internal/sessions/transaction.go @@ -0,0 +1,44 @@ +package sessions + +import ( + "context" + "encoding/json" +) + +// The procedures in this package run inside their caller's Session transaction, +// under the Session lock, and read and write only through the transaction +// interfaces declared beside them. Each interface lists exactly what its +// procedure calls. A procedure never begins, commits or rolls back a +// transaction: when it fails, the caller's transaction rolls back everything the +// procedure wrote. + +// LockedSession is what locking a Session row shows. +type LockedSession struct { + // Deleted reports that the Session was publicly deleted. Its row stays so + // that remaining execution can settle. + Deleted bool +} + +// Public applies the public view of a locked Session: a deleted Session is +// ErrNotFound. +func (s LockedSession) Public() error { + if s.Deleted { + return ErrNotFound + } + return nil +} + +// JournalTx reads the Session's public usage and journals its public changes. +type JournalTx interface { + // LoadUsage reads the Session's public usage as the transaction has left it. + LoadUsage(ctx context.Context) (json.RawMessage, error) + // AppendChanges journals public changes in order. + AppendChanges(ctx context.Context, changes ...SessionChange) error +} + +// ActiveTurnTx reads the Session's active Turn. +type ActiveTurnTx interface { + // LoadActiveTurn reads the Session's queued, in-progress or waiting Turn and + // reports whether it has one. + LoadActiveTurn(ctx context.Context) (Turn, bool, error) +} diff --git a/services/core/internal/sessions/transaction_test.go b/services/core/internal/sessions/transaction_test.go new file mode 100644 index 000000000..b1c29e1f8 --- /dev/null +++ b/services/core/internal/sessions/transaction_test.go @@ -0,0 +1,193 @@ +package sessions + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" +) + +// fakeTx is a strict Session transaction. Each method records its call, with +// the arguments that identify it, then runs its func; a method whose func is +// unset fails the test. Tests set only the funcs they expect and assert the +// exact call log. +type fakeTx struct { + t *testing.T + calls []string + + loadUsage func() (json.RawMessage, error) + appendChanges func([]SessionChange) error + loadActiveTurn func() (Turn, bool, error) + requestTurnCancel func() error + loadTurn func() (Turn, error) + loadEnding func() (Ending, error) + applyTurnEnd func(TurnEnd) error + cancelPendingInput func() error + failPendingInput func() error + loadEnvironmentInput func() (*EnvironmentInputState, error) + loadEnvironment func() (Environment, error) + recordEnvironmentFailure func() (time.Time, error) + expireEnvironment func() error + loadComputeSuspension func() (bool, error) + loadPendingFileWrite func() (bool, error) + loadBoundDevice func() (bool, error) + insertEnvironmentDevice func() error +} + +var ( + _ EnvironmentTerminationTx = (*fakeTx)(nil) + _ InputStartTx = (*fakeTx)(nil) + _ ComputeAdmissionTx = (*fakeTx)(nil) + _ EnvironmentDeviceTx = (*fakeTx)(nil) +) + +func (f *fakeTx) record(name string, set bool, detail ...string) { + f.t.Helper() + if !set { + f.t.Fatalf("unexpected call to %s", name) + } + f.calls = append(f.calls, strings.Join(append([]string{name}, detail...), " ")) +} + +func (f *fakeTx) LoadUsage(context.Context) (json.RawMessage, error) { + f.record("LoadUsage", f.loadUsage != nil) + return f.loadUsage() +} + +func (f *fakeTx) AppendChanges(_ context.Context, changes ...SessionChange) error { + f.record("AppendChanges", f.appendChanges != nil, strings.Join(types(changes), ",")) + return f.appendChanges(changes) +} + +func (f *fakeTx) LoadActiveTurn(context.Context) (Turn, bool, error) { + f.record("LoadActiveTurn", f.loadActiveTurn != nil) + return f.loadActiveTurn() +} + +func (f *fakeTx) RequestTurnCancel(_ context.Context, turn string) error { + f.record("RequestTurnCancel", f.requestTurnCancel != nil, turn) + return f.requestTurnCancel() +} + +func (f *fakeTx) LoadTurn(_ context.Context, turn string) (Turn, error) { + f.record("LoadTurn", f.loadTurn != nil, turn) + return f.loadTurn() +} + +func (f *fakeTx) LoadEnding(_ context.Context, turn string) (Ending, error) { + f.record("LoadEnding", f.loadEnding != nil, turn) + return f.loadEnding() +} + +func (f *fakeTx) ApplyTurnEnd(_ context.Context, turn string, end TurnEnd) error { + f.record("ApplyTurnEnd", f.applyTurnEnd != nil, turn) + return f.applyTurnEnd(end) +} + +func (f *fakeTx) CancelPendingInput(context.Context) error { + f.record("CancelPendingInput", f.cancelPendingInput != nil) + return f.cancelPendingInput() +} + +func (f *fakeTx) FailPendingInput(context.Context) error { + f.record("FailPendingInput", f.failPendingInput != nil) + return f.failPendingInput() +} + +func (f *fakeTx) LoadEnvironmentInput(context.Context) (*EnvironmentInputState, error) { + f.record("LoadEnvironmentInput", f.loadEnvironmentInput != nil) + return f.loadEnvironmentInput() +} + +func (f *fakeTx) LoadEnvironment(context.Context) (Environment, error) { + f.record("LoadEnvironment", f.loadEnvironment != nil) + return f.loadEnvironment() +} + +func (f *fakeTx) RecordEnvironmentFailure(_ context.Context, environment, reason string, _ *ProvisioningFailureDetail) (time.Time, error) { + f.record("RecordEnvironmentFailure", f.recordEnvironmentFailure != nil, environment, reason) + return f.recordEnvironmentFailure() +} + +func (f *fakeTx) ExpireEnvironment(_ context.Context, environment string) error { + f.record("ExpireEnvironment", f.expireEnvironment != nil, environment) + return f.expireEnvironment() +} + +func (f *fakeTx) LoadComputeSuspension(context.Context) (bool, error) { + f.record("LoadComputeSuspension", f.loadComputeSuspension != nil) + return f.loadComputeSuspension() +} + +func (f *fakeTx) LoadPendingFileWrite(context.Context) (bool, error) { + f.record("LoadPendingFileWrite", f.loadPendingFileWrite != nil) + return f.loadPendingFileWrite() +} + +func (f *fakeTx) LoadBoundDevice(context.Context) (bool, error) { + f.record("LoadBoundDevice", f.loadBoundDevice != nil) + return f.loadBoundDevice() +} + +func (f *fakeTx) InsertEnvironmentDevice(_ context.Context, device ExecutionDevice, credentialHash string) error { + f.record("InsertEnvironmentDevice", f.insertEnvironmentDevice != nil, device.ID, device.Name, device.EnvironmentID, credentialHash) + return f.insertEnvironmentDevice() +} + +// returns is a fake method that reads value. +func returns[T any](value T) func() (T, error) { + return func() (T, error) { return value, nil } +} + +// done is a fake method that applies successfully. +func done() error { return nil } + +// activeTurn is a fake LoadActiveTurn: the Session's active Turn, or none. +func activeTurn(turn *Turn) func() (Turn, bool, error) { + return func() (Turn, bool, error) { + if turn == nil { + return Turn{}, false, nil + } + return *turn, true, nil + } +} + +// inputs is a fake LoadEnvironmentInput that reads states one call at a time, +// then nil. +func inputs(states ...*EnvironmentInputState) func() (*EnvironmentInputState, error) { + return func() (*EnvironmentInputState, error) { + var state *EnvironmentInputState + if len(states) > 0 { + state, states = states[0], states[1:] + } + return state, nil + } +} + +// collect is a fake AppendChanges that keeps the journaled changes. +func collect(changes *[]SessionChange) func([]SessionChange) error { + return func(appended []SessionChange) error { + *changes = append(*changes, appended...) + return nil + } +} + +func assertCalls(t *testing.T, f *fakeTx, want ...string) { + t.Helper() + if strings.Join(f.calls, "\n") != strings.Join(want, "\n") { + t.Fatalf("calls:\n%s\nwant:\n%s", strings.Join(f.calls, "\n"), strings.Join(want, "\n")) + } +} + +var errStorage = errors.New("storage") + +func TestLockedSessionPublic(t *testing.T) { + if err := (LockedSession{}).Public(); err != nil { + t.Fatal(err) + } + if err := (LockedSession{Deleted: true}).Public(); !errors.Is(err, ErrNotFound) { + t.Fatal(err) + } +} diff --git a/services/core/internal/store/admin_session_archive.go b/services/core/internal/store/admin_session_archive.go index 6de7cae50..73f041409 100644 --- a/services/core/internal/store/admin_session_archive.go +++ b/services/core/internal/store/admin_session_archive.go @@ -10,6 +10,7 @@ import ( "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/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" @@ -63,7 +64,7 @@ func (s *Store) archiveManagedSession(ctx context.Context, tenantID, sessionID s if err != nil { return err } - kind, err := storedEnvironmentType(environment) + kind, err := sessions.EnvironmentType(environment.Configuration) if err != nil || kind != "openai_hosted" { return sessions.ErrInvalidInput } @@ -103,13 +104,14 @@ func (s *Store) archiveManagedSession(ctx context.Context, tenantID, sessionID s if allocated && allocation.RuntimeAllocation.State != "released" && allocation.RuntimeAllocation.ProviderKey != current.InstallationID { return deployment.ErrConflict } - if err := withEnvironmentInputActivity(ctx, q, session, func() error { + 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 cancelSessionWork(ctx, q, session) + return sessions.CancelWork(ctx, bound) }); err != nil { return err } diff --git a/services/core/internal/store/admin_summary.go b/services/core/internal/store/admin_summary.go index d9d665b5f..43364b8e6 100644 --- a/services/core/internal/store/admin_summary.go +++ b/services/core/internal/store/admin_summary.go @@ -6,6 +6,7 @@ import ( "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" @@ -48,7 +49,7 @@ func (s *Store) ReadAdminSummary(ctx context.Context, tenantID string, filter Ad return err } for _, row := range rows { - session, err := sessionFromRow(row.Session) + session, err := sessionpg.SessionFromRow(row.Session) if err != nil { return err } diff --git a/services/core/internal/store/artifact_capture.go b/services/core/internal/store/artifact_capture.go index 7a449027f..c27754f8c 100644 --- a/services/core/internal/store/artifact_capture.go +++ b/services/core/internal/store/artifact_capture.go @@ -4,13 +4,13 @@ import ( "archive/tar" "context" "encoding/json" - "errors" "io" "io/fs" "strings" "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" @@ -22,7 +22,7 @@ const MaxArtifactBatchBytes int64 = 500 << 20 // StageTurnArtifacts stores a complete export privately without locking admission during transfer. func (s *Store) StageTurnArtifacts(ctx context.Context, tenantID, sessionID, turnID, environmentID string, input io.Reader) error { - lookup, err := turnLookup(tenantID, sessionID, turnID) + lookup, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return err } @@ -53,15 +53,12 @@ func (s *Store) StageTurnArtifacts(ctx context.Context, tenantID, sessionID, tur return err } q := s.queries.WithTx(tx) - locked, err := q.LockSession(ctx, sqlc.LockSessionParams{TenantID: lookup.TenantID, ID: lookup.SessionID}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } + locked, err := sessionpg.LockSession(ctx, q, lookup.TenantID, lookup.SessionID) if err != nil { return err } - if locked.DeletedAt.Valid { - return sessions.ErrNotFound + if err := locked.Public(); err != nil { + return err } turn, err := q.GetTurn(ctx, lookup) if err != nil { diff --git a/services/core/internal/store/artifact_lifecycle.go b/services/core/internal/store/artifact_lifecycle.go index d0118232d..eb244e2fa 100644 --- a/services/core/internal/store/artifact_lifecycle.go +++ b/services/core/internal/store/artifact_lifecycle.go @@ -5,6 +5,7 @@ import ( "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" @@ -13,7 +14,7 @@ import ( // 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 := turnLookup(tenantID, sessionID, turnID) + lookup, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return err } diff --git a/services/core/internal/store/devices.go b/services/core/internal/store/devices.go index 9d0cf9827..e7866b42a 100644 --- a/services/core/internal/store/devices.go +++ b/services/core/internal/store/devices.go @@ -12,6 +12,7 @@ import ( "github.com/jackc/pgx/v5/pgtype" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -63,7 +64,7 @@ func (s *Store) GetDeviceCredential(ctx context.Context, deviceID string) (runti } func (s *Store) RevokeDevice(ctx context.Context, tenantID, deviceID string) error { - params, err := deviceLookup(tenantID, deviceID) + params, err := sessionpg.DeviceLookup(tenantID, deviceID) if err != nil { return err } @@ -77,7 +78,7 @@ func (s *Store) RevokeDevice(ctx context.Context, tenantID, deviceID string) err // BindSessionDevice keeps retries stable and refuses silent filesystem moves. // Dispatchers must obtain the full Session binding before delivery. func (s *Store) BindSessionDevice(ctx context.Context, tenantID, sessionID, deviceID string) error { - params, err := deviceLookup(tenantID, deviceID) + params, err := sessionpg.DeviceLookup(tenantID, deviceID) if err != nil { return err } @@ -96,7 +97,7 @@ func (s *Store) BindSessionDevice(ctx context.Context, tenantID, sessionID, devi } func (s *Store) GetSessionDevice(ctx context.Context, tenantID, sessionID string) (sessions.ExecutionDevice, error) { - params, err := deviceLookup(tenantID, sessionID) + params, err := sessionpg.DeviceLookup(tenantID, sessionID) if err != nil { return sessions.ExecutionDevice{}, err } @@ -109,7 +110,7 @@ func (s *Store) GetSessionDevice(ctx context.Context, tenantID, sessionID string // GetSessionRuntimeDevice reports an authorized connection binding. It does not // admit native execution or file access before Environment preparation completes. func (s *Store) GetSessionRuntimeDevice(ctx context.Context, tenantID, sessionID string) (sessions.ExecutionDevice, error) { - params, err := deviceLookup(tenantID, sessionID) + params, err := sessionpg.DeviceLookup(tenantID, sessionID) if err != nil { return sessions.ExecutionDevice{}, err } @@ -124,7 +125,7 @@ func (s *Store) GetSessionRuntimeDevice(ctx context.Context, tenantID, sessionID } func (s *Store) GetSessionExecutionBinding(ctx context.Context, tenantID, sessionID string) (sessions.ExecutionBinding, error) { - params, err := deviceLookup(tenantID, sessionID) + params, err := sessionpg.DeviceLookup(tenantID, sessionID) if err != nil { return sessions.ExecutionBinding{}, err } @@ -153,16 +154,6 @@ func executionDevice(id pgtype.UUID, name string, environmentID pgtype.UUID) ses return value } -func deviceLookup(tenantID, id string) (sqlc.GetDeviceParams, error) { - var p sqlc.GetDeviceParams - var err error - if p.TenantID, err = parseID(tenantID); err != nil { - return p, err - } - p.ID, err = parseID(id) - return p, err -} - func (s *Store) TouchRuntimeHeartbeat(ctx context.Context, deviceID string) (runtimedevice.HeartbeatStatus, error) { id, err := parseID(deviceID) if err != nil { diff --git a/services/core/internal/store/environment_connections.go b/services/core/internal/store/environment_connections.go index 941b53085..ed2b4f247 100644 --- a/services/core/internal/store/environment_connections.go +++ b/services/core/internal/store/environment_connections.go @@ -2,10 +2,8 @@ package store import ( "context" - "encoding/json" "errors" - 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/persistence/postgres/sessionpg" @@ -100,12 +98,13 @@ func (s *Store) withEnvironmentConnection(ctx context.Context, tenant, environme if err != nil { return err } - return withEnvironmentInputActivity(ctx, q, session, func() error { return apply(ctx, q, row) }) + return sessions.TrackInputActivity(ctx, sessionpg.BindSession(q, tenantID, session), func(ctx context.Context) error { return apply(ctx, q, row) }) }) } func recordEnvironmentConnection(ctx context.Context, q *sqlc.Queries, row sqlc.GetSessionEnvironmentRow, status string) error { - if _, err := storedEnvironmentType(row); err != nil { + kind, err := sessions.EnvironmentType(row.Configuration) + if err != nil { return err } if status != "connected" && status != "disconnected" { @@ -114,35 +113,7 @@ func recordEnvironmentConnection(ctx context.Context, q *sqlc.Queries, row sqlc. if err := q.SetEnvironmentConnectionStatus(ctx, sqlc.SetEnvironmentConnectionStatusParams{ID: row.Environment.ID, Status: status}); err != nil { return err } - return recordEnvironmentState(ctx, q, row, status) -} - -// recordEnvironmentState appends the pinned Environment state event for an -// already committed status. A failure uses the observed official error; the -// failed step travels only in the separate error event and Session error. -func recordEnvironmentState(ctx context.Context, q *sqlc.Queries, row sqlc.GetSessionEnvironmentRow, status string) error { - kind, err := storedEnvironmentType(row) - if err != nil { - return err - } - state := &v1.SessionEnvironmentState{ID: uuid.UUID(row.Environment.ID.Bytes).String(), Type: kind, Status: status} - if status == "failed" { - state.Error = &v1.StreamError{Type: "environment_error", Code: "environment_connection_failed", Message: "The environment failed to connect."} - } - return sessionpg.AppendChanges(ctx, q, row.Environment.SessionID, sessions.SessionChange{Event: v1.SessionEvent{ - Type: "agent.session.environment." + status, - Environment: state, - }}) -} - -func storedEnvironmentType(row sqlc.GetSessionEnvironmentRow) (string, error) { - var config struct { - Type string `json:"type"` - } - if err := json.Unmarshal(row.Configuration, &config); err != nil || (config.Type != "self_hosted" && config.Type != "openai_hosted") { - return "", errors.New("invalid stored Environment type") - } - return config.Type, nil + return sessionpg.AppendChanges(ctx, q, row.Environment.SessionID, sessions.EnvironmentStateChange(uuid.UUID(row.Environment.ID.Bytes).String(), kind, status)) } func parseConnectionGeneration(value string) (pgtype.UUID, error) { diff --git a/services/core/internal/store/environment_file_writes.go b/services/core/internal/store/environment_file_writes.go index 14ec762bf..0c56b14f3 100644 --- a/services/core/internal/store/environment_file_writes.go +++ b/services/core/internal/store/environment_file_writes.go @@ -11,6 +11,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/auditpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/writeaudit" ) @@ -73,10 +74,11 @@ func (s *Store) ReserveEnvironmentFileWrite(ctx context.Context, tenant, environ if uuid.UUID(device.ID.Bytes).String() != key.DeviceID || device.EnvironmentID != lookup.EnvironmentID { return sessions.ErrDeviceBindingConflict } - if err := checkRuntimeComputeAdmission(ctx, q, session); err != nil { + bound := sessionpg.BindSession(q, lookup.TenantID, session) + if err := sessions.CheckComputeAdmission(ctx, bound); err != nil { return err } - if err := environmentInputMayStart(ctx, q, session); err != nil { + if err := sessions.CheckInputStart(ctx, bound); err != nil { return err } pending, err := q.EnvironmentFileWriteHasPendingInput(ctx, session) @@ -190,11 +192,3 @@ func fileWriteFromRow(row sqlc.EnvironmentFileWrite, session pgtype.UUID) sessio } return result } - -func checkEnvironmentFileWriteGate(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - blocked, err := q.EnvironmentFileWriteBlocksSession(ctx, session) - if err == nil && blocked { - return sessions.ErrTurnConflict - } - return err -} diff --git a/services/core/internal/store/environment_initialization.go b/services/core/internal/store/environment_initialization.go index b5cb9bf6b..6bc15cbaf 100644 --- a/services/core/internal/store/environment_initialization.go +++ b/services/core/internal/store/environment_initialization.go @@ -5,6 +5,7 @@ import ( "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" @@ -106,6 +107,10 @@ func (s *Store) FailEnvironmentInitialization(ctx context.Context, owner session if err := q.FailEnvironmentInitialization(ctx, row.Environment.ID); err != nil { return err } - return failEnvironment(ctx, q, row, row.Environment.SessionID, failure.Reason(), failure.Detail(), func() error { return cancelSessionWork(ctx, q, row.Environment.SessionID) }) + environment, err := sessionpg.EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, nil) + if err != nil { + return err + } + return sessions.FailEnvironment(ctx, sessionpg.BindSession(q, row.TenantID, row.Environment.SessionID), environment, failure.Reason(), failure.Detail()) }) } diff --git a/services/core/internal/store/environment_input_activity.go b/services/core/internal/store/environment_input_activity.go index a16e516e9..5e32b7e0e 100644 --- a/services/core/internal/store/environment_input_activity.go +++ b/services/core/internal/store/environment_input_activity.go @@ -2,90 +2,22 @@ package store import ( "context" - "errors" - v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" + "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" - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" ) -func environmentInputActivity(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) (*sessions.EnvironmentInputActivity, error) { - activity, _, err := environmentInputState(ctx, q, session) - return activity, err -} - -// environmentInputState also reports whether the latest reservation is still -// pending and can start a Turn, including a provisioning hosted initial input -// that has no public activity. While a Turn is active or newer than it, the -// reservation is not reported. -func environmentInputState(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) (*sessions.EnvironmentInputActivity, bool, error) { - row, err := q.GetEnvironmentInputActivity(ctx, session) - if errors.Is(err, pgx.ErrNoRows) { - return nil, false, nil - } - if err != nil { - return nil, false, err - } - pending := row.State == sessions.EnvironmentInputPending - activity := &sessions.EnvironmentInputActivity{Status: "idle", LastActiveAt: row.CreatedAt.Time} - if row.SettledAt.Valid { - activity.LastActiveAt = row.SettledAt.Time - } - if row.State == sessions.EnvironmentInputFailed { - activity.Status, activity.Failure = "failed", "environment_unavailable" - if row.FailureCode.Valid { - activity.Failure = row.FailureCode.String - } - } - if row.IsInitial && row.State == sessions.EnvironmentInputExpired { - activity.Status = "failed" - } - if row.EnvironmentType == "openai_hosted" && row.IsInitial && - (row.State == sessions.EnvironmentInputPending || row.State == sessions.EnvironmentInputCancelled) { - // No Turn has started. The pinned Session contract permits idle while a - // hosted Environment provisions; neither a caller action nor an invented - // in-progress/idle transition is appropriate here. - return nil, pending, nil - } - if pending && row.EnvironmentType != "openai_hosted" && row.ConnectionStatus != "connected" { - activity.Status = "requires_action" - activity.EnvironmentID = uuid.UUID(row.EnvironmentID.Bytes).String() - } - return activity, pending, nil -} - -func withEnvironmentInputActivity(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, apply func() error) error { - before, err := environmentInputActivity(ctx, q, session) +// 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 } - if err := apply(); err != nil { - return err - } - after, pending, err := environmentInputState(ctx, q, session) - if err != nil || after == nil { - // Admitted input is represented by the normal Turn and Session events. - return err - } - if before != nil && before.Status == after.Status && before.EnvironmentID == after.EnvironmentID && before.Failure == after.Failure { - return nil - } - usage, err := sessionpg.LoadUsage(ctx, q, session) - if err != nil { - return err - } - return sessionpg.AppendChanges(ctx, q, session, sessions.SessionChange{ - Event: v1.SessionEvent{Type: "agent.session." + after.Status}, - EnvironmentInputActivity: after, SessionUsage: usage, Settled: !pending, - }) -} - -func (s *Store) withEnvironmentInputSession(ctx context.Context, tenant, session string, apply func(context.Context, *sqlc.Queries, pgtype.UUID) error) error { - return s.withPublicSession(ctx, tenant, session, func(ctx context.Context, q *sqlc.Queries, id pgtype.UUID) error { - return withEnvironmentInputActivity(ctx, q, id, func() error { return apply(ctx, q, id) }) + 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_expiry.go b/services/core/internal/store/environment_input_expiry.go index 7578c3be4..3724937da 100644 --- a/services/core/internal/store/environment_input_expiry.go +++ b/services/core/internal/store/environment_input_expiry.go @@ -5,6 +5,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/jackc/pgx/v5" ) @@ -22,7 +23,7 @@ func (s *Store) ExpireEnvironmentInputs(ctx context.Context) (int64, error) { return err } for _, row := range rows { - err := withEnvironmentInputActivity(ctx, q, row.SessionID, func() error { + 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 { diff --git a/services/core/internal/store/environment_inputs.go b/services/core/internal/store/environment_inputs.go index c41a2cf56..5fbfedec1 100644 --- a/services/core/internal/store/environment_inputs.go +++ b/services/core/internal/store/environment_inputs.go @@ -12,6 +12,7 @@ import ( "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" ) @@ -59,7 +60,7 @@ func (s *Store) ReserveEnvironmentInput(ctx context.Context, tenantID, sessionID result = sessions.EnvironmentInputReservation{SessionID: sessionID, State: sessions.EnvironmentInputAdmitted, Receipts: receipts} return audit() } - if err := checkEnvironmentFileWriteGate(ctx, q, session); err != nil { + 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}) @@ -69,7 +70,7 @@ func (s *Store) ReserveEnvironmentInput(ctx context.Context, tenantID, sessionID return err } if environment.Environment.Status == "failed" { - if kind, err := storedEnvironmentType(environment); err == nil && kind == "openai_hosted" { + if kind, err := sessions.EnvironmentType(environment.Configuration); err == nil && kind == "openai_hosted" { return sessions.ErrHostedEnvironmentFailed } } @@ -211,7 +212,11 @@ func settleEnvironmentInput(ctx context.Context, q *sqlc.Queries, tenantID strin return sessions.EnvironmentInputReservation{}, err } if state == sessions.EnvironmentInputAdmitted { - if err := environmentInputMayStart(ctx, q, row.SessionID); err != nil { + 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 { @@ -227,7 +232,7 @@ func settleEnvironmentInput(ctx context.Context, q *sqlc.Queries, tenantID strin return sessions.EnvironmentInputReservation{}, err } if state == sessions.EnvironmentInputAdmitted { - params, err := turnLookup(tenantID, result.SessionID, result.Receipts[0].TurnID) + params, err := sessionpg.TurnLookup(tenantID, result.SessionID, result.Receipts[0].TurnID) if err != nil { return sessions.EnvironmentInputReservation{}, err } @@ -242,20 +247,6 @@ func settleEnvironmentInput(ctx context.Context, q *sqlc.Queries, tenantID strin return result, nil } -func environmentInputMayStart(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - if err := checkEnvironmentFileWriteGate(ctx, q, session); err != nil { - return err - } - _, err := q.GetActiveTurn(ctx, session) - if errors.Is(err, pgx.ErrNoRows) { - return nil - } - if err == nil { - return sessions.ErrTurnConflict - } - return err -} - 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 { diff --git a/services/core/internal/store/environment_setup.go b/services/core/internal/store/environment_setup.go index bd4a25ddd..6ba986f7b 100644 --- a/services/core/internal/store/environment_setup.go +++ b/services/core/internal/store/environment_setup.go @@ -8,6 +8,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -56,7 +57,7 @@ func (s *Store) saveEnvironmentSetup(ctx context.Context, q *sqlc.Queries, tenan func (s *Store) ReadEnvironmentSetup(ctx context.Context, tenant, session string) (environmentconfig.Setup, error) { var result environmentconfig.Setup - lookup, err := deviceLookup(tenant, session) + lookup, err := sessionpg.DeviceLookup(tenant, session) if err != nil { return result, sessions.ErrNotFound } diff --git a/services/core/internal/store/environments.go b/services/core/internal/store/environments.go index eb206c0a4..5907788ac 100644 --- a/services/core/internal/store/environments.go +++ b/services/core/internal/store/environments.go @@ -3,16 +3,14 @@ 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/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" ) @@ -45,7 +43,7 @@ func (s *Store) GetEnvironment(ctx context.Context, tenantID, environmentID stri } id := pgunit.PathID(environmentID) row, err := s.queries.GetEnvironment(ctx, sqlc.GetEnvironmentParams{TenantID: tenant, ID: id}) - return environmentFromRow(row.Environment, row.TenantID, row.Configuration, err) + return sessionpg.EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, err) } func (s *Store) GetSessionEnvironment(ctx context.Context, tenantID, sessionID string) (sessions.Environment, error) { @@ -58,23 +56,5 @@ func (s *Store) GetSessionEnvironment(ctx context.Context, tenantID, sessionID s return sessions.Environment{}, err } row, err := s.queries.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: tenant, ID: id}) - return environmentFromRow(row.Environment, row.TenantID, row.Configuration, err) -} - -func environmentFromRow(row sqlc.Environment, tenant pgtype.UUID, configuration []byte, err error) (sessions.Environment, error) { - if errors.Is(err, pgx.ErrNoRows) { - return sessions.Environment{}, sessions.ErrNotFound - } - if err != nil { - return sessions.Environment{}, fmt.Errorf("get environment: %w", err) - } - configuration, err = jsonobject.Normalize(configuration) - if err != nil { - return sessions.Environment{}, fmt.Errorf("decode environment configuration: %w: %w", sessions.ErrInvalidInput, err) - } - return sessions.Environment{ - ID: uuid.UUID(row.ID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), - TenantID: uuid.UUID(tenant.Bytes).String(), Status: row.Status, - Initialization: row.Initialization, CreatedAt: row.CreatedAt.Time, Configuration: configuration, - }, nil + return sessionpg.EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, err) } diff --git a/services/core/internal/store/execution.go b/services/core/internal/store/execution.go index c820c55b2..05238c038 100644 --- a/services/core/internal/store/execution.go +++ b/services/core/internal/store/execution.go @@ -4,19 +4,12 @@ import ( "context" "errors" - "github.com/jackc/pgx/v5" - "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") -// transactor runs one transaction: pgunit's Pool or its execution Lease. -type transactor interface { - Transaction(context.Context, func(context.Context, pgx.Tx) error) error -} - // NewExecution returns the execution writer built on lease, which the caller // acquired and closes. The writer's Session and execution-only transactions run // on the leased connection; reads keep the pool. Keep public admission on the diff --git a/services/core/internal/store/function_calls.go b/services/core/internal/store/function_calls.go index 32c8a13d9..95fe59a82 100644 --- a/services/core/internal/store/function_calls.go +++ b/services/core/internal/store/function_calls.go @@ -7,6 +7,7 @@ import ( "strings" "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" @@ -14,7 +15,7 @@ import ( // RecordFunctionCall commits an execution callback and its required-action state together. func (s *Store) RecordFunctionCall(ctx context.Context, tenantID, sessionID, turnID string, call sessions.FunctionCall) error { - p, err := turnLookup(tenantID, sessionID, turnID) + p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return err } @@ -54,7 +55,7 @@ func (s *Store) RecordFunctionCall(ctx context.Context, tenantID, sessionID, tur } func (s *Store) GetFunctionCall(ctx context.Context, tenantID, sessionID, turnID, callID string) (sessions.FunctionCall, error) { - p, err := turnLookup(tenantID, sessionID, turnID) + p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return sessions.FunctionCall{}, err } @@ -73,7 +74,7 @@ func (s *Store) GetFunctionCall(ctx context.Context, tenantID, sessionID, turnID // PendingFunctionCalls excludes applied results and cancelling or terminal Turns. func (s *Store) PendingFunctionCalls(ctx context.Context, tenantID, sessionID, turnID string) ([]sessions.FunctionCall, error) { - p, err := turnLookup(tenantID, sessionID, turnID) + p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return nil, err } diff --git a/services/core/internal/store/function_results.go b/services/core/internal/store/function_results.go index 55283b705..ebe4ace4f 100644 --- a/services/core/internal/store/function_results.go +++ b/services/core/internal/store/function_results.go @@ -8,6 +8,7 @@ import ( "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/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" @@ -44,7 +45,7 @@ func (s *Store) ConfirmFunctionResult(ctx context.Context, tenantID, sessionID, } func (s *Store) withFunctionCall(ctx context.Context, tenantID, sessionID, turnID, callID string, fn func(context.Context, *sqlc.Queries, sqlc.Turn, sqlc.FunctionCall) error) error { - p, err := turnLookup(tenantID, sessionID, turnID) + p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return err } diff --git a/services/core/internal/store/function_state.go b/services/core/internal/store/function_state.go index 0423e07c9..c557deefd 100644 --- a/services/core/internal/store/function_state.go +++ b/services/core/internal/store/function_state.go @@ -40,7 +40,7 @@ func recordFunctionState(ctx context.Context, q *sqlc.Queries, turn sqlc.Turn) e if err != nil { return err } - if err := sessionpg.AppendChanges(ctx, q, turn.SessionID, sessions.TurnChanges(turnFromRow(turn), false)...); err != nil { + if err := sessionpg.AppendChanges(ctx, q, turn.SessionID, sessions.TurnChanges(sessionpg.TurnFromRow(turn), false)...); err != nil { return err } } @@ -48,5 +48,5 @@ func recordFunctionState(ctx context.Context, q *sqlc.Queries, turn sqlc.Turn) e if err != nil { return err } - return sessionpg.AppendChanges(ctx, q, turn.SessionID, sessions.ActivityChange(turnFromRow(turn), usage, actions)) + return sessionpg.AppendChanges(ctx, q, turn.SessionID, sessions.ActivityChange(sessionpg.TurnFromRow(turn), usage, actions)) } diff --git a/services/core/internal/store/initial_files.go b/services/core/internal/store/initial_files.go index a8a51b732..9b8e7f2da 100644 --- a/services/core/internal/store/initial_files.go +++ b/services/core/internal/store/initial_files.go @@ -10,6 +10,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/files" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -84,7 +85,7 @@ func readInitialSourceFile(ctx context.Context, tx pgx.Tx, source sqlc.SourceFil // ReadInitialEnvironmentFile decrypts only the next frozen file, bounding memory per installation. func (s *Store) ReadInitialEnvironmentFile(ctx context.Context, tenant, session string, position int) (environmentconfig.InitialFileMetadata, []byte, error) { - lookup, err := deviceLookup(tenant, session) + lookup, err := sessionpg.DeviceLookup(tenant, session) if err != nil { return environmentconfig.InitialFileMetadata{}, nil, err } diff --git a/services/core/internal/store/local_environment_devices.go b/services/core/internal/store/local_environment_devices.go index 92691f6c0..684dbfeec 100644 --- a/services/core/internal/store/local_environment_devices.go +++ b/services/core/internal/store/local_environment_devices.go @@ -2,13 +2,11 @@ package store import ( "context" - "encoding/json" - "errors" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/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" ) @@ -18,13 +16,10 @@ func (s *Store) CreateEnvironmentDevice(ctx context.Context, tenantID, environme if err != nil { return sessions.ExecutionDevice{}, err } - var configuration struct { - Type string `json:"type"` - } - if json.Unmarshal(environment.Configuration, &configuration) != nil || configuration.Type != "openai_hosted" { + if kind, err := sessions.EnvironmentType(environment.Configuration); err != nil || kind != "openai_hosted" { return sessions.ExecutionDevice{}, sessions.ErrInvalidInput } - lookup, err := deviceLookup(tenantID, environment.ID) + lookup, err := sessionpg.DeviceLookup(tenantID, environment.ID) if err != nil { return sessions.ExecutionDevice{}, err } @@ -32,33 +27,12 @@ func (s *Store) CreateEnvironmentDevice(ctx context.Context, tenantID, environme if err != nil { return sessions.ExecutionDevice{}, err } + device := sessions.ExecutionDevice{ID: uuid.UUID(params.ID.Bytes).String(), Name: params.Name, EnvironmentID: environment.ID} err = s.withPublicSession(ctx, tenantID, environment.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - return createEnvironmentDevice(ctx, q, lookup, session, params) + return sessions.CreateEnvironmentDevice(ctx, sessionpg.BindSession(q, lookup.TenantID, session), device, params.CredentialHash.String) }) if err != nil { return sessions.ExecutionDevice{}, err } - return sessions.ExecutionDevice{ID: uuid.UUID(params.ID.Bytes).String(), Name: params.Name, EnvironmentID: environment.ID}, nil -} - -func createEnvironmentDevice(ctx context.Context, q *sqlc.Queries, lookup sqlc.GetDeviceParams, session pgtype.UUID, params sqlc.CreateDeviceParams) error { - _, err := q.GetSessionDevice(ctx, sqlc.GetSessionDeviceParams{TenantID: lookup.TenantID, ID: session}) - if err == nil { - return sessions.ErrDeviceBindingConflict - } - if !errors.Is(err, pgx.ErrNoRows) { - return err - } - id, err := q.CreateEnvironmentDevice(ctx, sqlc.CreateEnvironmentDeviceParams{ - ID: params.ID, TenantID: params.TenantID, Name: params.Name, - CredentialHash: params.CredentialHash, EnvironmentID: lookup.ID, - }) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrDeviceBindingConflict - } - if err != nil { - return err - } - _, err = q.BindSessionDevice(ctx, sqlc.BindSessionDeviceParams{TenantID: lookup.TenantID, ID: session, ID_2: id}) - return err + return device, nil } diff --git a/services/core/internal/store/runtime_allocation_state.go b/services/core/internal/store/runtime_allocation_state.go index 71059f182..669e9446e 100644 --- a/services/core/internal/store/runtime_allocation_state.go +++ b/services/core/internal/store/runtime_allocation_state.go @@ -8,6 +8,7 @@ import ( "github.com/jackc/pgx/v5/pgtype" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -64,11 +65,11 @@ func (s *Store) requestRuntimeCleanup(ctx context.Context, owner RuntimeAllocati if err != nil { return sqlc.RuntimeAllocation{}, err } - cancel := func() error { return cancelSessionWork(ctx, q, current.SessionID) } + session := sessionpg.BindSession(q, current.TenantID, current.SessionID) if current.DeletedAt.Valid { - err = cancel() + err = sessions.CancelWork(ctx, session) } else { - err = terminateRuntimeEnvironment(ctx, q, current, reason, detail, cancel) + err = sessions.TerminateEnvironment(ctx, session, current.Expired, reason, detail) } if err != nil { return sqlc.RuntimeAllocation{}, err @@ -121,7 +122,7 @@ func (s *Store) mutateRuntimeAllocation(ctx context.Context, owner RuntimeAlloca if previous.ID != owner.ID || previous.DeviceID != owner.DeviceID || previous.ProviderKey != owner.ProviderKey || previous.NodeID != owner.NodeID { return RuntimeAllocation{}, sessions.ErrIdempotencyConflict } - lookup, _ := deviceLookup(owner.TenantID, owner.EnvironmentID) + lookup, _ := sessionpg.DeviceLookup(owner.TenantID, owner.EnvironmentID) var result RuntimeAllocation err = s.withSession(ctx, owner.TenantID, previous.SessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { current, err := q.GetRuntimeAllocation(ctx, sqlc.GetRuntimeAllocationParams{TenantID: lookup.TenantID, EnvironmentID: lookup.ID}) diff --git a/services/core/internal/store/runtime_allocations.go b/services/core/internal/store/runtime_allocations.go index 5cbd323e3..5a9e45fa9 100644 --- a/services/core/internal/store/runtime_allocations.go +++ b/services/core/internal/store/runtime_allocations.go @@ -12,6 +12,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -65,7 +66,7 @@ func (s *Store) ReserveRuntimeAllocation(ctx context.Context, tenant, environmen if json.Unmarshal(owned.Configuration, &config) != nil || config.Type != "openai_hosted" { return RuntimeAllocation{}, sessions.ErrInvalidInput } - lookup, err := deviceLookup(tenant, environment) + lookup, err := sessionpg.DeviceLookup(tenant, environment) if err != nil { return RuntimeAllocation{}, err } @@ -114,7 +115,8 @@ func (s *Store) ReserveRuntimeAllocation(ctx context.Context, tenant, environmen nodeID = placement.NodeID generation = placement.DeploymentGeneration } - if err := createEnvironmentDevice(ctx, q, lookup, session, device); err != nil { + dedicated := sessions.ExecutionDevice{ID: uuid.UUID(device.ID.Bytes).String(), Name: device.Name, EnvironmentID: owned.ID} + if err := sessions.CreateEnvironmentDevice(ctx, sessionpg.BindSession(q, lookup.TenantID, session), dedicated, device.CredentialHash.String); err != nil { return err } row, err := q.CreateRuntimeAllocation(ctx, sqlc.CreateRuntimeAllocationParams{ @@ -134,7 +136,7 @@ func (s *Store) ReserveRuntimeAllocation(ctx context.Context, tenant, environmen // GetRuntimeAllocation is an internal cleanup lookup, including deleted Sessions. func (s *Store) GetRuntimeAllocation(ctx context.Context, tenant, environment string) (RuntimeAllocation, error) { - lookup, err := deviceLookup(tenant, environment) + lookup, err := sessionpg.DeviceLookup(tenant, environment) if err != nil { return RuntimeAllocation{}, err } diff --git a/services/core/internal/store/runtime_enrollment.go b/services/core/internal/store/runtime_enrollment.go index dbed9d76c..7ec1dbdd7 100644 --- a/services/core/internal/store/runtime_enrollment.go +++ b/services/core/internal/store/runtime_enrollment.go @@ -5,6 +5,7 @@ import ( "errors" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -30,7 +31,7 @@ func (s *Store) EnrollRuntime(ctx context.Context, environmentID, credentialHash if err != nil { return RuntimeEnrollment{}, err } - lookup, err := deviceLookup(tenant, environmentID) + lookup, err := sessionpg.DeviceLookup(tenant, environmentID) if err != nil { return RuntimeEnrollment{}, err } diff --git a/services/core/internal/store/runtime_environment_terminal.go b/services/core/internal/store/runtime_environment_terminal.go deleted file mode 100644 index 86e7de6d4..000000000 --- a/services/core/internal/store/runtime_environment_terminal.go +++ /dev/null @@ -1,97 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - - 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/sessionpg" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5/pgtype" -) - -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 -} - -// terminateRuntimeEnvironment participates in the allocation's Session transaction. -// Public expiry does not assert compute removal or invent an expired SSE variant. -// A first failure records the hosted provisioning failure; an already terminal -// Environment only settles remaining input, without repeating events. -func terminateRuntimeEnvironment(ctx context.Context, q *sqlc.Queries, current sqlc.GetRuntimeAllocationRow, reason string, detail *sessions.ProvisioningFailureDetail, cancel func() error) error { - row, err := q.GetSessionEnvironment(ctx, sqlc.GetSessionEnvironmentParams{TenantID: current.TenantID, ID: current.SessionID}) - if err != nil { - return err - } - terminal := row.Environment.Status == "expired" || row.Environment.Status == "failed" - if !terminal && !current.Expired { - return failEnvironment(ctx, q, row, current.SessionID, reason, detail, cancel) - } - return withEnvironmentInputActivity(ctx, q, current.SessionID, func() error { - if !terminal { - if err := q.SetEnvironmentConnectionStatus(ctx, sqlc.SetEnvironmentConnectionStatusParams{ID: row.Environment.ID, Status: "expired"}); err != nil { - return err - } - } - if err := q.FailSessionEnvironmentInput(ctx, current.SessionID); err != nil { - return err - } - return cancel() - }) -} - -// failEnvironment records, in the caller's transaction and in the observed -// official order, agent.session.environment.failed, an error event carrying the -// safe reason, then one agent.session.failed snapshot. The snapshot captures the -// settled input activity, Usage and the failure, matching later Session reads. -// Pending input settles as failed exactly as before. -func failEnvironment(ctx context.Context, q *sqlc.Queries, row sqlc.GetSessionEnvironmentRow, session pgtype.UUID, reason string, detail *sessions.ProvisioningFailureDetail, cancel func() error) error { - var rawDetail []byte - if detail != nil { - var err error - rawDetail, err = json.Marshal(detail) - if err != nil { - return err - } - } - failedAt, err := q.RecordEnvironmentFailure(ctx, sqlc.RecordEnvironmentFailureParams{ID: row.Environment.ID, FailureReason: pgtype.Text{String: reason, Valid: true}, FailureDetail: rawDetail}) - if err != nil { - return err - } - if err := recordEnvironmentState(ctx, q, row, "failed"); err != nil { - return err - } - if err := q.FailSessionEnvironmentInput(ctx, session); err != nil { - return err - } - if err := cancel(); err != nil { - return err - } - activity, pending, err := environmentInputState(ctx, q, session) - if err != nil { - return err - } - usage, err := sessionpg.LoadUsage(ctx, q, session) - if err != nil { - return err - } - if err := sessionpg.AppendChanges(ctx, q, session, sessions.SessionChange{Event: v1.SessionEvent{ - Type: "error", Error: &v1.StreamError{Type: "environment_error", Code: "sandbox_error", Message: reason}, - }}); err != nil { - return err - } - return sessionpg.AppendChanges(ctx, q, session, sessions.SessionChange{ - Event: v1.SessionEvent{Type: "agent.session.failed"}, - EnvironmentInputActivity: activity, EnvironmentFailure: &sessions.EnvironmentFailure{Reason: reason, FailedAt: failedAt.Time}, - SessionUsage: usage, Settled: !pending, - }) -} diff --git a/services/core/internal/store/runtime_lifecycle_nodes.go b/services/core/internal/store/runtime_lifecycle_nodes.go index c931cdf05..a387dfc2b 100644 --- a/services/core/internal/store/runtime_lifecycle_nodes.go +++ b/services/core/internal/store/runtime_lifecycle_nodes.go @@ -8,6 +8,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" @@ -96,7 +97,7 @@ func (s *Store) ListUnallocatedHostedEnvironmentsForNode(ctx context.Context, no // ResolveRuntimeLifecycleNode routes direct provisioning before an allocation // exists. An existing allocation must agree with its immutable placement. func (s *Store) ResolveRuntimeLifecycleNode(ctx context.Context, tenant, environment string) (string, error) { - lookup, err := deviceLookup(tenant, environment) + lookup, err := sessionpg.DeviceLookup(tenant, environment) if err != nil { return "", err } diff --git a/services/core/internal/store/runtime_suspension.go b/services/core/internal/store/runtime_suspension.go index 90d4afff3..9b0e2a5f6 100644 --- a/services/core/internal/store/runtime_suspension.go +++ b/services/core/internal/store/runtime_suspension.go @@ -7,6 +7,7 @@ import ( "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" @@ -116,7 +117,7 @@ func (s *Store) RuntimeActivity(ctx context.Context, owner RuntimeAllocation) (R // TouchRuntimeActivity is used only by operations requiring live compute. // Public history and published-artifact reads do not call it. func (s *Store) TouchRuntimeActivity(ctx context.Context, tenant, environment string) error { - lookup, err := deviceLookup(tenant, environment) + lookup, err := sessionpg.DeviceLookup(tenant, environment) if err != nil { return err } @@ -171,13 +172,3 @@ func (s *Store) CountRuntimeRetainedAllocations(ctx context.Context, provider st } return s.queries.CountRuntimeRetainedAllocations(ctx, id) } - -// checkRuntimeComputeAdmission runs under the Session lock before a new durable -// execution owner is created. Existing receipts remain readable in every phase. -func checkRuntimeComputeAdmission(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - blocked, err := q.RuntimeComputeBlocksAdmission(ctx, session) - if err == nil && blocked { - return sessions.ErrTurnConflict - } - return err -} diff --git a/services/core/internal/store/runtime_suspension_concurrency_test.go b/services/core/internal/store/runtime_suspension_concurrency_test.go index 1e32c700b..09f8d9f9a 100644 --- a/services/core/internal/store/runtime_suspension_concurrency_test.go +++ b/services/core/internal/store/runtime_suspension_concurrency_test.go @@ -8,6 +8,7 @@ import ( "testing" "time" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -199,7 +200,7 @@ func TestRuntimeSuspensionQuiesceCannotOvertakeClaim(t *testing.T) { if _, err := tx.Exec(ctx, `INSERT INTO turns(id,session_id,status) VALUES($1,$2,'queued')`, turn, owner.SessionID); err != nil { t.Fatal(err) } - params, err := turnLookup(owner.TenantID, owner.SessionID, turn) + params, err := sessionpg.TurnLookup(owner.TenantID, owner.SessionID, turn) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/store/scheduling.go b/services/core/internal/store/scheduling.go index 863e527c4..032aba557 100644 --- a/services/core/internal/store/scheduling.go +++ b/services/core/internal/store/scheduling.go @@ -2,10 +2,12 @@ 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" @@ -111,7 +113,7 @@ func (s *Store) SessionStreamSnapshot(ctx context.Context, tenantID, sessionID s if err != nil { return err } - if session, err = sessionFromRow(row); err != nil { + if session, err = sessionpg.SessionFromRow(row); err != nil { return err } cursor = row.EventSequence @@ -134,16 +136,17 @@ func readSessionActivity(ctx context.Context, q *sqlc.Queries, session sessions. 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) + 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) - session.EnvironmentInputActivity, session.PendingInput, err = environmentInputState(ctx, q, id) + 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 } @@ -154,7 +157,7 @@ func readSessionActivity(ctx context.Context, q *sqlc.Queries, session sessions. if err != nil { return session, err } - turn := turnFromRow(row) + turn := sessionpg.TurnFromRow(row) session.LastTurn = &turn session.RequiredActions, err = functionActions(ctx, q, row) if err != nil { @@ -163,3 +166,15 @@ func readSessionActivity(ctx context.Context, q *sqlc.Queries, session sessions. 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/session_artifacts.go b/services/core/internal/store/session_artifacts.go index 6bacd5e86..d0c880770 100644 --- a/services/core/internal/store/session_artifacts.go +++ b/services/core/internal/store/session_artifacts.go @@ -8,6 +8,7 @@ import ( "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/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -111,13 +112,13 @@ func (s *Store) DeleteSessionArtifact(ctx context.Context, tenantID, sessionID, return s.pooled.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { q := s.queries.WithTx(tx) // Use the same lock order as whole-Session deletion and Turn publication. - locked, err := q.LockSession(ctx, sqlc.LockSessionParams{TenantID: lookup.TenantID, ID: lookup.SessionID}) - if errors.Is(err, pgx.ErrNoRows) || err == nil && locked.DeletedAt.Valid { - return sessions.ErrNotFound - } + locked, err := sessionpg.LockSession(ctx, q, lookup.TenantID, lookup.SessionID) if err != nil { return err } + if err := locked.Public(); err != nil { + return err + } oid, err := q.DeleteSessionArtifact(ctx, sqlc.DeleteSessionArtifactParams(lookup)) if errors.Is(err, pgx.ErrNoRows) { return sessions.ErrNotFound diff --git a/services/core/internal/store/session_creation_identity.go b/services/core/internal/store/session_creation_identity.go index 225728b2c..6f59d0555 100644 --- a/services/core/internal/store/session_creation_identity.go +++ b/services/core/internal/store/session_creation_identity.go @@ -15,6 +15,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" + "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" @@ -154,7 +155,7 @@ func (s *Store) FindSessionCreation(ctx context.Context, tenantID, key string, r if row.CreationRequestHash.String != hash.String { return sessions.Creation{}, sessions.ErrIdempotencyConflict } - session, err := sessionFromRow(row) + session, err := sessionpg.SessionFromRow(row) session.Environment = environment // The row and cursor share one committed snapshot; later events remain observable. return sessions.Creation{Session: session, Cursor: row.EventSequence}, err diff --git a/services/core/internal/store/session_creator.go b/services/core/internal/store/session_creator.go deleted file mode 100644 index 8b2f8f78b..000000000 --- a/services/core/internal/store/session_creator.go +++ /dev/null @@ -1,19 +0,0 @@ -package store - -import ( - "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" - "github.com/jackc/pgx/v5/pgtype" -) - -func sessionCreator(kind, id pgtype.Text) (*identity.Subject, error) { - if !kind.Valid && !id.Valid { - return nil, nil - } - creator := identity.Subject{Kind: kind.String, ID: id.String} - if !kind.Valid || !id.Valid || creator.Validate() != nil { - return nil, fmt.Errorf("invalid stored Session creator") - } - return &creator, nil -} diff --git a/services/core/internal/store/session_deletion.go b/services/core/internal/store/session_deletion.go index bc074f526..7f6772abd 100644 --- a/services/core/internal/store/session_deletion.go +++ b/services/core/internal/store/session_deletion.go @@ -20,23 +20,23 @@ import ( // 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 sqlc.LockSessionRow) 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.ID.Bytes).String(), "") + return auditpg.RecordWriteAudit(ctx, q, tenantID, "delete", "session", uuid.UUID(session.Bytes).String(), "") } - if session.DeletedAt.Valid { + if locked.Deleted { return audit() } - if err := requireSessionSettled(ctx, q, session.ID); err != nil { + if err := requireSessionSettled(ctx, q, session); err != nil { return err } - if err := q.DeleteSessionArtifacts(ctx, session.ID); err != nil { + if err := q.DeleteSessionArtifacts(ctx, session); err != nil { return err } - if err := q.ReleaseUnallocatedRuntimePlacement(ctx, session.ID); err != nil { + if err := q.ReleaseUnallocatedRuntimePlacement(ctx, session); err != nil { return err } - if err := q.MarkSessionDeleted(ctx, session.ID); err != nil { + if err := q.MarkSessionDeleted(ctx, session); err != nil { return err } return audit() @@ -53,61 +53,12 @@ func requireSessionSettled(ctx context.Context, q *sqlc.Queries, session pgtype. } else if !errors.Is(err, pgx.ErrNoRows) { return err } - _, pending, err := environmentInputState(ctx, q, session) + state, err := sessionpg.LoadEnvironmentInput(ctx, q, session) if err != nil { return err } - if pending { + if _, pending := sessions.InputActivity(state); pending { return sessions.ErrNotIdle } return nil } - -// cancelSessionWork requests cancellation of active work for Runtime cleanup. -func cancelSessionWork(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - turn, err := q.GetActiveTurn(ctx, session) - if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return err - } - if err == nil { - if err := requestTurnCancel(ctx, q, session, turn); err != nil { - return err - } - } - if err := q.CancelSessionEnvironmentInput(ctx, session); err != nil { - return err - } - return nil -} - -func requestTurnCancel(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, turn sqlc.Turn) error { - if err := q.RequestTurnCancel(ctx, sqlc.RequestTurnCancelParams{ID: turn.ID, SessionID: session}); err != nil { - return err - } - if turn.Status == sessions.TurnQueued { - cancelled, err := q.SessionEventTurn(ctx, sqlc.SessionEventTurnParams{SessionID: session, ID: turn.ID}) - if err != nil { - return err - } - ending, err := sessionpg.LoadEnding(ctx, q, session, cancelled.ID) - if err != nil { - return err - } - if err := sessionpg.ApplyTurnEnd(ctx, q, session, cancelled.ID, sessions.EndTurn(turnFromRow(cancelled), ending)); err != nil { - return err - } - } else if turn.Status == sessions.TurnWaiting && !turn.CancelRequestedAt.Valid { - cancelling, err := q.SessionEventTurn(ctx, sqlc.SessionEventTurnParams{SessionID: session, ID: turn.ID}) - if err != nil { - return err - } - usage, err := sessionpg.LoadUsage(ctx, q, session) - if err != nil { - return err - } - if err := sessionpg.AppendChanges(ctx, q, session, sessions.ActivityChange(turnFromRow(cancelling), usage, nil)); err != nil { - return err - } - } - return nil -} diff --git a/services/core/internal/store/session_deletion_test.go b/services/core/internal/store/session_deletion_test.go index 28d25bb6d..c2e3668c8 100644 --- a/services/core/internal/store/session_deletion_test.go +++ b/services/core/internal/store/session_deletion_test.go @@ -10,6 +10,7 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -22,8 +23,12 @@ import ( // the marker together. Upgraded databases can retain such markers, so hidden // work must still settle; tests use this to reach that state. func (s *Store) commitLegacyDeletion(ctx context.Context, tenantID, sessionID string) error { + tenant, err := parseID(tenantID) + if err != nil { + return err + } return s.withPublicSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - if err := cancelSessionWork(ctx, q, session); err != nil { + if err := sessions.CancelWork(ctx, sessionpg.BindSession(q, tenant, session)); err != nil { return err } if err := q.DeleteSessionArtifacts(ctx, session); err != nil { diff --git a/services/core/internal/store/session_diagnostics.go b/services/core/internal/store/session_diagnostics.go index 1b435c539..316c65a21 100644 --- a/services/core/internal/store/session_diagnostics.go +++ b/services/core/internal/store/session_diagnostics.go @@ -6,6 +6,7 @@ import ( "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" @@ -25,7 +26,7 @@ func (s *Store) GetSessionDiagnosticsSnapshot(ctx context.Context, tenantID, ses if err != nil { return err } - session, err = sessionFromRow(row) + session, err = sessionpg.SessionFromRow(row) if err != nil { return err } @@ -56,11 +57,11 @@ func (s *Store) GetTurnDiagnosticsSnapshot(ctx context.Context, tenantID, sessio if err != nil { return err } - result.Session, err = sessionFromRow(row) + result.Session, err = sessionpg.SessionFromRow(row) if err != nil { return err } - result.Turn = turnFromRow(turn) + result.Turn = sessionpg.TurnFromRow(turn) rows, err := q.ListTurnItemDiagnostics(ctx, sqlc.ListTurnItemDiagnosticsParams{SessionID: params.SessionID, TurnID: params.ID}) if err != nil { return err diff --git a/services/core/internal/store/session_environment_snapshot.go b/services/core/internal/store/session_environment_snapshot.go index 334aebaaf..315fe08a9 100644 --- a/services/core/internal/store/session_environment_snapshot.go +++ b/services/core/internal/store/session_environment_snapshot.go @@ -5,6 +5,7 @@ import ( "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" ) @@ -14,7 +15,7 @@ func sessionEnvironmentSnapshot(ctx context.Context, q *sqlc.Queries, session sq if errors.Is(err, pgx.ErrNoRows) { return nil, nil } - value, err := environmentFromRow(row.Environment, row.TenantID, row.Configuration, err) + value, err := sessionpg.EnvironmentFromRow(row.Environment, row.TenantID, row.Configuration, err) if err != nil { return nil, err } diff --git a/services/core/internal/store/session_initial_input.go b/services/core/internal/store/session_initial_input.go index 878fb7ea2..18f34a14b 100644 --- a/services/core/internal/store/session_initial_input.go +++ b/services/core/internal/store/session_initial_input.go @@ -129,7 +129,7 @@ func (s *Store) createSessionResources(ctx context.Context, tenant string, param key := uuid.NewString() if environment != nil { // Environment input waits for the existing preparation and leased promotion. - if err := withEnvironmentInputActivity(ctx, q, row.ID, func() error { + if err := sessions.TrackInputActivity(ctx, sessionpg.BindSession(q, row.TenantID, row.ID), func(ctx context.Context) error { _, err := q.CreateEnvironmentInputReservation(ctx, sqlc.CreateEnvironmentInputReservationParams{ ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, SessionID: row.ID, IdempotencyKey: key, Batch: encodedInput, IsInitial: true, diff --git a/services/core/internal/store/session_metadata.go b/services/core/internal/store/session_metadata.go index f96a87e53..24ae3c2ff 100644 --- a/services/core/internal/store/session_metadata.go +++ b/services/core/internal/store/session_metadata.go @@ -9,6 +9,7 @@ import ( "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" @@ -40,6 +41,6 @@ func (s *Store) UpdateSessionMetadata(ctx context.Context, tenantID, sessionID s if err != nil { return sessions.Session{}, fmt.Errorf("update session metadata: %w", err) } - session, decodeErr := sessionFromRow(row) + session, decodeErr := sessionpg.SessionFromRow(row) return s.sessionActivity(ctx, session, decodeErr) } diff --git a/services/core/internal/store/session_transaction.go b/services/core/internal/store/session_transaction.go index 1f14909e2..ef0749ffb 100644 --- a/services/core/internal/store/session_transaction.go +++ b/services/core/internal/store/session_transaction.go @@ -2,9 +2,7 @@ package store import ( "context" - "errors" - "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" @@ -13,8 +11,8 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -// All Turn admission and lifecycle writes lock the tenant-scoped Session first. -// This orders inputs against completion/cancellation across service processes. +// All Turn admission and lifecycle writes run in sessionpg's Session +// transaction on the writer. func (s *Store) withSession(ctx context.Context, tenantID, sessionID string, apply func(context.Context, *sqlc.Queries, pgtype.UUID) error) error { return s.withSessionState(ctx, tenantID, sessionID, false, apply) } @@ -24,17 +22,19 @@ 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 sqlc.LockSessionRow) error { - if public && session.DeletedAt.Valid { - return sessions.ErrNotFound + 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.ID) + return apply(ctx, q, session) }) } -// withLockedSession locks the tenant-owned Session row, including a publicly -// deleted one, and commits only when apply succeeds. -func (s *Store) withLockedSession(ctx context.Context, tenantID, sessionID string, public bool, apply func(context.Context, *sqlc.Queries, sqlc.LockSessionRow) error) error { +// 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 @@ -46,17 +46,7 @@ func (s *Store) withLockedSession(ctx context.Context, tenantID, sessionID strin return err } } - return s.writer.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { - q := s.queries.WithTx(tx) - session, err := q.LockSession(ctx, sqlc.LockSessionParams{TenantID: tenant, ID: id}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } else if err != nil { - return err - } - if err := apply(ctx, q, session); err != nil { - return err - } - return sessionpg.PruneChanges(ctx, q, id) + 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) }) } diff --git a/services/core/internal/store/sessions.go b/services/core/internal/store/sessions.go index ac1a97acb..03fc02b09 100644 --- a/services/core/internal/store/sessions.go +++ b/services/core/internal/store/sessions.go @@ -22,6 +22,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/metadata" "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" ) @@ -34,7 +35,7 @@ type Store struct { // writer runs Session and execution-only transactions, and lease grants // execution authority. New sets writer to pooled and leaves lease nil; // NewExecution sets both to the lease it borrows. Neither changes later. - writer transactor + writer pgunit.Transactor lease *pgunit.Lease credentialCipher *credentialcrypto.Cipher // publicURL is OAC_PUBLIC_URL. Core derives every address it gives @@ -156,7 +157,7 @@ func (s *Store) createSession(ctx context.Context, tenantID string, input sessio if err != nil { return sessions.Creation{}, fmt.Errorf("create session: %w", err) } - session, err := sessionFromRow(row) + session, err := sessionpg.SessionFromRow(row) session.Environment = environment return sessions.Creation{Session: session, Created: row.ID == params.ID, Cursor: row.EventSequence}, err } @@ -175,7 +176,7 @@ func (s *Store) GetSession(ctx context.Context, tenantID, sessionID string) (ses if err != nil { return sessions.Session{}, fmt.Errorf("get session: %w", err) } - session, decodeErr := sessionFromRow(row) + session, decodeErr := sessionpg.SessionFromRow(row) return s.sessionActivity(ctx, session, decodeErr) } @@ -211,7 +212,7 @@ func (s *Store) ListSessions(ctx context.Context, tenantID, cursor string, limit rows = rows[:limit] } for _, row := range rows { - session, err := sessionFromRow(row) + session, err := sessionpg.SessionFromRow(row) session, err = s.sessionActivity(ctx, session, err) if err != nil { return sessions.Page{}, err @@ -231,21 +232,3 @@ func parseID(value string) (pgtype.UUID, error) { } return id, nil } - -func sessionFromRow(row sqlc.Session) (sessions.Session, error) { - session := sessions.Session{ID: uuid.UUID(row.ID.Bytes).String(), TenantID: uuid.UUID(row.TenantID.Bytes).String(), Engine: row.Engine, CreatedAt: row.CreatedAt.Time, RequiredActions: []v1.FunctionCallAction{}} - creator, err := sessionCreator(row.CreatorKind, row.CreatorID) - if err != nil { - return sessions.Session{}, err - } - session.Creator = creator - configuration, err := jsonobject.Normalize(row.Configuration) - if err != nil { - return sessions.Session{}, fmt.Errorf("decode session configuration: %w: %w", sessions.ErrInvalidInput, err) - } - session.Configuration = configuration - if err := json.Unmarshal(row.Metadata, &session.Metadata); err != nil { - return sessions.Session{}, fmt.Errorf("decode session metadata: %w", err) - } - return session, nil -} diff --git a/services/core/internal/store/subagent_identities.go b/services/core/internal/store/subagent_identities.go index e234046f4..2d66352cc 100644 --- a/services/core/internal/store/subagent_identities.go +++ b/services/core/internal/store/subagent_identities.go @@ -9,6 +9,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -50,7 +51,7 @@ func projectSubagentIdentity(ctx context.Context, q *sqlc.Queries, session, turn // GetSubagentIdentity recovers a binding in its authorized, visible Session. func (s *Store) GetSubagentIdentity(ctx context.Context, tenantID, sessionID, nativeID string) (SubagentIdentity, error) { - p, err := deviceLookup(tenantID, sessionID) + p, err := sessionpg.DeviceLookup(tenantID, sessionID) if err != nil { return SubagentIdentity{}, err } diff --git a/services/core/internal/store/turn_completion.go b/services/core/internal/store/turn_completion.go index 731476f62..64f76d487 100644 --- a/services/core/internal/store/turn_completion.go +++ b/services/core/internal/store/turn_completion.go @@ -17,7 +17,7 @@ import ( // 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 := turnLookup(tenantID, sessionID, turnID) + p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return sessions.Turn{}, err } @@ -95,10 +95,10 @@ func (s *Store) CompleteExecution(ctx context.Context, tenantID, sessionID, turn if err != nil { return err } - return sessionpg.ApplyTurnEnd(ctx, q, session, row.ID, sessions.EndTurn(turnFromRow(row), ending)) + return sessionpg.ApplyTurnEnd(ctx, q, session, row.ID, sessions.EndTurn(sessionpg.TurnFromRow(row), ending)) }) if err != nil { return sessions.Turn{}, err } - return turnFromRow(row), nil + return sessionpg.TurnFromRow(row), nil } diff --git a/services/core/internal/store/turn_events.go b/services/core/internal/store/turn_events.go index bb9ef93bd..30d0ab773 100644 --- a/services/core/internal/store/turn_events.go +++ b/services/core/internal/store/turn_events.go @@ -9,6 +9,7 @@ import ( "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/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" @@ -16,7 +17,7 @@ import ( // AppendTurnEvents records an ordered batch atomically, not public SSE replay events. func (s *Store) AppendTurnEvents(ctx context.Context, tenantID, sessionID, turnID string, first int32, events []sessions.ExecutionEvent) error { - p, err := turnLookup(tenantID, sessionID, turnID) + p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return err } @@ -96,7 +97,7 @@ func insertTurnEvent(ctx context.Context, q *sqlc.Queries, turn sqlc.Turn, kind } func (s *Store) ListTurnEvents(ctx context.Context, tenantID, sessionID, turnID string, after int32, limit int) ([]sessions.TurnEvent, error) { - p, err := turnLookup(tenantID, sessionID, turnID) + p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return nil, err } diff --git a/services/core/internal/store/turn_inputs.go b/services/core/internal/store/turn_inputs.go index 8c731d98b..919b4f5a4 100644 --- a/services/core/internal/store/turn_inputs.go +++ b/services/core/internal/store/turn_inputs.go @@ -46,6 +46,10 @@ func (s *Store) SubmitInputs(ctx context.Context, tenantID, sessionID, key strin 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) @@ -57,7 +61,7 @@ func (s *Store) SubmitInputs(ctx context.Context, tenantID, sessionID, key strin 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 := checkEnvironmentFileWriteGate(ctx, q, session); err != nil { + if err := sessions.CheckFileWriteGate(ctx, sessionpg.BindSession(q, tenant, session)); err != nil { return err } } @@ -134,7 +138,7 @@ func admitInput(ctx context.Context, q *sqlc.Queries, tenantID string, session p 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(turnFromRow(turn), true)...) + err = sessionpg.AppendChanges(ctx, q, session, sessions.TurnChanges(sessionpg.TurnFromRow(turn), true)...) } } else { err = nil // Retain even an idle cancellation's retry identity. @@ -150,7 +154,11 @@ func admitInput(ctx context.Context, q *sqlc.Queries, tenantID string, session p return sessions.InputReceipt{}, err } if input.Kind == "cancel" && turn.ID.Valid { - if err := requestTurnCancel(ctx, q, session, turn); err != nil { + tenant, err := parseID(tenantID) + if err != nil { + return sessions.InputReceipt{}, err + } + if err := sessions.CancelTurn(ctx, sessionpg.BindSession(q, tenant, session), sessionpg.TurnFromRow(turn)); err != nil { return sessions.InputReceipt{}, err } } @@ -164,7 +172,7 @@ func admitInput(ctx context.Context, q *sqlc.Queries, tenantID string, session p if err != nil { return sessions.InputReceipt{}, err } - if err := sessionpg.AppendChanges(ctx, q, session, sessions.ActivityChange(turnFromRow(turn), usage, nil)); err != nil { + if err := sessionpg.AppendChanges(ctx, q, session, sessions.ActivityChange(sessionpg.TurnFromRow(turn), usage, nil)); err != nil { return sessions.InputReceipt{}, err } } @@ -173,7 +181,7 @@ func admitInput(ctx context.Context, q *sqlc.Queries, tenantID string, session p // 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 := turnLookup(tenantID, sessionID, turnID) + params, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return nil, err } diff --git a/services/core/internal/store/turn_reads.go b/services/core/internal/store/turn_reads.go index 0f9d06655..cdff33971 100644 --- a/services/core/internal/store/turn_reads.go +++ b/services/core/internal/store/turn_reads.go @@ -6,6 +6,7 @@ import ( "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" @@ -42,7 +43,7 @@ func (s *Store) ListTurns(ctx context.Context, tenantID, sessionID, cursor strin rows = rows[:limit] } for _, row := range rows { - page.Turns = append(page.Turns, turnFromRow(row)) + page.Turns = append(page.Turns, sessionpg.TurnFromRow(row)) } return page, nil } diff --git a/services/core/internal/store/turns.go b/services/core/internal/store/turns.go index db40430a4..2023153ff 100644 --- a/services/core/internal/store/turns.go +++ b/services/core/internal/store/turns.go @@ -2,11 +2,9 @@ package store import ( "context" - "encoding/json" "errors" "fmt" - "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" @@ -31,14 +29,14 @@ func (s *Store) GetTurn(ctx context.Context, tenantID, sessionID, turnID string) if err != nil { return sessions.Turn{}, fmt.Errorf("get turn: %w", err) } - return turnFromRow(row), nil + 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 := turnLookup(tenantID, sessionID, turnID) + params, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) if err != nil { return sessions.Turn{}, err } @@ -62,7 +60,7 @@ func (s *Store) TransitionTurn(ctx context.Context, tenantID, sessionID, turnID if err != nil { return sessions.Turn{}, fmt.Errorf("transition turn: %w", err) } - return turnFromRow(row), nil + return sessionpg.TurnFromRow(row), nil } func transitionTurn(ctx context.Context, q *sqlc.Queries, params sqlc.GetTurnParams, input sessions.TurnTransition) (sqlc.Turn, error) { @@ -72,7 +70,7 @@ func transitionTurn(ctx context.Context, q *sqlc.Queries, params sqlc.GetTurnPar return sqlc.Turn{}, err } if input.ExpectedStatus == sessions.TurnQueued && input.Status == sessions.TurnInProgress { - if err := checkRuntimeComputeAdmission(ctx, q, params.SessionID); err != nil { + if err := sessions.CheckComputeAdmission(ctx, sessionpg.BindSession(q, params.TenantID, params.SessionID)); err != nil { return sqlc.Turn{}, err } } @@ -87,7 +85,7 @@ func transitionTurn(ctx context.Context, q *sqlc.Queries, params sqlc.GetTurnPar return sqlc.Turn{}, err } if !sessions.TerminalStatus(row.Status) { - if err := sessionpg.AppendChanges(ctx, q, row.SessionID, sessions.TurnChanges(turnFromRow(row), false)...); err != nil { + if err := sessionpg.AppendChanges(ctx, q, row.SessionID, sessions.TurnChanges(sessionpg.TurnFromRow(row), false)...); err != nil { return sqlc.Turn{}, err } return row, nil @@ -102,7 +100,7 @@ func transitionTurn(ctx context.Context, q *sqlc.Queries, params sqlc.GetTurnPar if err != nil { return sqlc.Turn{}, err } - if err := sessionpg.ApplyTurnEnd(ctx, q, row.SessionID, row.ID, sessions.EndTurn(turnFromRow(row), ending)); err != nil { + 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 @@ -121,31 +119,9 @@ func validTransition(from, to string) bool { } } -func turnLookup(tenantID, sessionID, turnID string) (sqlc.GetTurnParams, error) { - var p sqlc.GetTurnParams - var err error - if p.TenantID, err = parseID(tenantID); err != nil { - return p, err - } - if p.SessionID, err = parseID(sessionID); err != nil { - return p, err - } - p.ID, err = parseID(turnID) - return p, err -} - // 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 } - -func turnFromRow(row sqlc.Turn) sessions.Turn { - return sessions.Turn{ - ID: uuid.UUID(row.ID.Bytes).String(), SessionID: uuid.UUID(row.SessionID.Bytes).String(), Status: row.Status, - CreatedAt: row.CreatedAt.Time, StartedAt: row.StartedAt.Time, CompletedAt: row.CompletedAt.Time, - CancelRequestedAt: row.CancelRequestedAt.Time, Outcome: json.RawMessage(row.Outcome), Usage: json.RawMessage(row.TokenUsage), - ArtifactCaptureStarted: row.ArtifactCaptureStarted, - } -}