diff --git a/services/core/IMPLEMENTATION.md b/services/core/IMPLEMENTATION.md index 22c31b5b7..2ad6b0b22 100644 --- a/services/core/IMPLEMENTATION.md +++ b/services/core/IMPLEMENTATION.md @@ -16,13 +16,13 @@ Two errors are shared across domains, each with one `api` helper: `textvalue.Err Shared vocabulary has one owner each, and domains use it rather than copy it. `internal/environmentconfig` owns Environment setup, Skills, Plugins and initial files with their validation and public metadata; `Setup.Validate` checks requested configuration, where a Skill may be an unresolved reference, and `Setup.ValidateInstalled` checks frozen, installable configuration. `internal/skills` owns `ParseVersion`, the canonical positive decimal Skill version. `internal/metadata` owns the metadata rules: `Validate` for the pair, key and value limits and U+0000, `ValidateStorable` for U+0000 alone, and `Encode` with its 64 KiB bound. `internal/jsonobject` owns `Normalize`, the stable encoding of stored JSON objects that snapshots and retry identities compare. These packages import no persistence. -`internal/sessions` owns the Session vocabulary: Sessions, Turns, inputs, Environments and their provisioning failures, function calls, Item and Artifact reads, executor credentials, and the errors Session operations return, which `api` maps in `writeSessionsError`. It also owns the Session change vocabulary and its decisions: the public changes that report Turn and Session transitions, what a Turn that ends settles, measured Turn usage and the Session activity a change reports. Sequences that several Session operations share are `sessions` procedures, each over the transaction interface it declares: `CancelTurn` and `CancelWork` cancel work, `TrackInputActivity` reports the input activity a write changes, `FailEnvironment` and `TerminateEnvironment` settle an Environment that failed or whose managed compute ended, `CreateEnvironmentDevice` creates and binds a hosted Environment's dedicated Runtime device, 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. +`internal/sessions` owns the Session vocabulary: Sessions, Turns, inputs, Environments and their provisioning failures, function calls, Item and Artifact reads, executor credentials, and the errors Session operations return, which `api` maps in `writeSessionsError`. It also owns the Session change vocabulary and its decisions: the public changes that report Turn and Session transitions, what a Turn that ends settles, measured Turn usage and the Session activity a change reports. Sequences that several Session operations share are `sessions` procedures, each over the transaction interface it declares: `CancelTurn` and `CancelWork` cancel work, `TrackInputActivity` reports the input activity a write changes, `FailEnvironment` and `TerminateEnvironment` settle an Environment that failed or whose managed compute ended, `CreateEnvironmentDevice` creates and binds a hosted Environment's dedicated Runtime device, `CheckComputeAdmission`, `CheckFileWriteGate` and `CheckInputStart` gate new work, `AdmitFunctionResult` saves a submitted function result for its Turn's call, and `LoadRequiredActions` lists the calls a Turn waits on. `sessions.ExecutionOperations` runs the Session writes only the execution owner makes; each function-call operation loads the Turn and its call in one leased Session transaction, decides in `sessions`, then applies in `sessionpg`. `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`. -`cmd/server` owns the execution lease. It acquires one `pgunit.Lease`, builds every lease-bound adapter on it, and passes the lease and those adapters together as one `execution.Owner` to `execution.StartWorker`. If anything fails before that call, `cmd/server` closes the lease. From that call the Worker owns cleanup: a failed start closes the lease before it returns, and a started Worker closes it after `Run` has cancelled and drained its work. Each close runs under its own bounded deadline, independent of the cancelled request or run. Lease-bound adapters and the store writer borrow the lease and never close it, and the Worker uses the lease only through `Owner.Lease`, never through an adapter. Store integration tests start the Worker the same way through `startWorker`. +`cmd/server` owns the execution lease. It acquires one `pgunit.Lease`, builds every lease-bound adapter on it, and passes the lease and those adapters together as one `execution.Owner` to `execution.StartWorker`, which binds the Dispatcher to the Owner's execution operations through `Dispatcher.Bind`. If anything fails before that call, `cmd/server` closes the lease. From that call the Worker owns cleanup: a failed start closes the lease before it returns, and a started Worker closes it after `Run` has cancelled and drained its work. Each close runs under its own bounded deadline, independent of the cancelled request or run. Lease-bound adapters and the store writer borrow the lease and never close it, and the Worker uses the lease only through `Owner.Lease`, never through an adapter. Store integration tests start the Worker the same way through `startWorker`, and bind a Dispatcher that runs Turns without a Worker through `Dispatcher.Bind`. Domain owners, each with its PostgreSQL adapter under `internal/persistence/postgres`: @@ -34,7 +34,7 @@ Domain owners, each with its PostgreSQL adapter under `internal/persistence/post - `modelconfiguration` (`modelconfigurationpg`): each Harness's deployment default model configuration and its last-use observations. - `skills` (`skillpg`): Skills and their immutable versions: archive checks, the default and latest pointers, version selection and deletion, and each version's sealed archive. Session creation freezes selected versions inside its `store` transaction with the `skills` rules. - `deployment` (`deploymentpg`): the sandbox deployment and its nodes: provider configuration and the sealed credential, the specification and retained generations, setup, update and switch, node enrollment, identity and authentication, generation configuration, capacity, presence, status, host history, reset, and the counts of nodes and sandboxes bound to the public address. It interprets Sandbox Provider declarations through the `providers.Registry` it is given, whose lookups return typed errors. `cmd/server` builds that registry and calls it only to build a direct Provider and to discover a Provider's configuration; it takes each setup's mode and declared operations from the `Setup` that `deployment` returns. The only other reader is `store`, whose own `providers.Builtin()` validates the stored specification at Session admission and reads the public-origin requirement at placement. Deployment changes run through `deployment.ExecutionOperations` on `deploymentpg.NewExecution(lease, …)`, which the Worker receives as `execution.Owner.Deployment`; the Worker reads the deployment and prepares a selection's setup through the pooled `deployment.Service` in `execution.Dispatcher.Deployment` and lists the Sessions a reset still has to archive through the pooled `deployment.Reader` in `execution.Dispatcher.DeploymentReader`, and node management and reads use the pooled `deploymentpg.Store`, which `cmd/server` also reads the owner epoch from and runs the host-history sampler on. Provider calls run outside transactions, and the final transaction rechecks the expected generation. Allocations, placement, reset admission and the reset's Session archive stay in `store`. -- `sessions` (`sessionpg`): Session use cases and reads as they leave `store`. The pooled `sessions.Service` runs the use cases on `sessionpg.Store`, which implements `sessions.Storage`, and plain reads use `sessions.Reader`, which `sessionpg.Store` also implements, directly. So far these cover root Items and Subagents (their reads) and Session Artifacts (their reads, deletion and the staging of a Turn's export). `cmd/server` builds one `sessionpg.Store` and wires the Service and the Store into `execution.Dispatcher.Sessions` and `SessionsReader` and into the api fields; the Worker stages Artifacts through the Service. +- `sessions` (`sessionpg`): Session use cases and reads as they leave `store`. The pooled `sessions.Service` runs the use cases on `sessionpg.Store`, which implements `sessions.Storage`, and plain reads use `sessions.Reader`, which `sessionpg.Store` also implements, directly. So far these cover root Items and Subagents (their reads) and Session Artifacts (their reads, deletion and the staging of a Turn's export). `cmd/server` builds one `sessionpg.Store` and wires the Service and the Store into `execution.Dispatcher.Sessions` and `SessionsReader` and into the api fields; the Worker stages Artifacts through the Service. Function calls and their application receipts run through `sessions.ExecutionOperations` on `sessionpg.NewExecution(lease)`, which the Worker receives as `execution.Owner.Sessions`. ## Request handling @@ -157,7 +157,7 @@ Core replaces a Turn's usage with each complete valid token breakdown, which `se Subagent observations use the common types in `internal/agentdaemon/proto/subagents.go`. Core assigns public IDs and projects them under the Session lock and the leased execution journal; native names, history parsing and outcome proof stay in the adapters. Child Turns have their own table, separate from Core's queue, and Session Turn reads and the Session event stream carry root work only. The migration-defined `public_execution_turns` view has no public reader; do not reintroduce mixed Session Turn pages. Repeated effects are idempotent, root output freezes first, and child Items are delivered before their terminal Turn snapshot. -Function calls are stored per tenant, Session and Turn with immutable public and executor call identities and arguments. Result admission and application receipts serialize on the Session lock with cancellation and terminal transitions. The store keeps the complete caller-validated result, including omitted versus null error and output and ordered text and image parts; identical retries return the saved decision and changed results conflict. Recording a call moves the Turn to `waiting`, and the last application receipt resumes it. A function result may join a message or cancel batch: its explicit Turn and call identity select an existing call, admission never creates a Turn for it, and the result and its input retry record commit in one transaction. Targets resolve only after the tenant Session lookup; an unknown call or a call of another Turn is 400 `invalid_request_error`. The input cursor skips function results, whose own receipts decide application. A Session with functions waits for a device that declares `function_tools`, and Core delivers each saved result once per live dispatch, appending a non-null error as a final text part because the native result has no error field. +Function calls are stored per tenant, Session and Turn with immutable public and executor call identities and arguments. Result admission and application receipts serialize on the Session lock with cancellation and terminal transitions. `sessions.AdmitFunctionResult` keeps the complete caller-validated result, including omitted versus null error and output and ordered text and image parts; identical retries return the saved decision and changed results conflict. Recording a call through `sessions.ExecutionOperations` moves the Turn to `waiting`, and the last application receipt resumes it. A function result may join a message or cancel batch: its explicit Turn and call identity select an existing call, admission never creates a Turn for it, and the result and its input retry record commit in one transaction. Targets resolve only after the tenant Session lookup; an unknown call or a call of another Turn is 400 `invalid_request_error`. The input cursor skips function results, whose own receipts decide application. A Session with functions waits for a device that declares `function_tools`, and Core delivers each saved result once per live dispatch, appending a non-null error as a final text part because the native result has no error field. Neutral tool and message observations go into the journal before Core projects public Items; the Item projector validates the shared observation contract and never decodes engine-native snapshots. `command_output` fragments update Items only for a command already indexed in the same Turn, and commit with their `agent.output.command_execution_output.delta` events. Completion output replaces accumulated drafts, terminal Items ignore late fragments, and cancellation keeps partial output without inventing a successful completion. diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index d3a39beaf..dad7ba5b5 100644 --- a/services/core/cmd/server/main.go +++ b/services/core/cmd/server/main.go @@ -301,11 +301,16 @@ func run() error { if err != nil { return errors.Join(err, lease.Close(ctx)) } + sessionExecution, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) + if err != nil { + return errors.Join(err, lease.Close(ctx)) + } // From this call on the Worker closes the lease, even when it fails to start. worker, err = execution.StartWorker(ctx, dispatcher, execution.Owner{ Lease: lease, Store: store.NewExecution(executionStore, lease), Deployment: deploymentExecution, + Sessions: sessionExecution, }) if err != nil { return err diff --git a/services/core/internal/execution/delivery.go b/services/core/internal/execution/delivery.go index df15ca24f..e0b901b15 100644 --- a/services/core/internal/execution/delivery.go +++ b/services/core/internal/execution/delivery.go @@ -105,7 +105,7 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe var pending *pendingInput var cancelSent time.Time var cancelReply <-chan cancellationResult - functions := &functionExchange{kind: request.AgentKind, store: d.Store, tenant: tenantID, session: sessionID, turn: request.RunID, tools: request.FunctionTools} + functions := &functionExchange{kind: request.AgentKind, store: d.Store, sessions: d.sessionExecution, tenant: tenantID, session: sessionID, turn: request.RunID, tools: request.FunctionTools} done := false cancelCtx, stopCancellation := context.WithCancel(ctx) defer stopCancellation() diff --git a/services/core/internal/execution/deployment_provider_observations_test.go b/services/core/internal/execution/deployment_provider_observations_test.go index 80c6c50df..2816644ea 100644 --- a/services/core/internal/execution/deployment_provider_observations_test.go +++ b/services/core/internal/execution/deployment_provider_observations_test.go @@ -73,7 +73,11 @@ func newFinishObservationFixture(t *testing.T, maxConnections int32) finishObser if err != nil { t.Fatal(err) } - return finishObservationFixture{s, owner.Store, owner.Lease, pool, defaults, tenant, session, Dispatcher{Store: owner.Store, Observer: defaults}} + dispatcher, err := (&Dispatcher{Observer: defaults}).Bind(owner) + if err != nil { + t.Fatal(err) + } + return finishObservationFixture{s, owner.Store, owner.Lease, pool, defaults, tenant, session, *dispatcher} } func (f finishObservationFixture) start(t *testing.T) sessions.InputReceipt { t.Helper() diff --git a/services/core/internal/execution/dispatcher.go b/services/core/internal/execution/dispatcher.go index fb99e4556..54ad6433d 100644 --- a/services/core/internal/execution/dispatcher.go +++ b/services/core/internal/execution/dispatcher.go @@ -29,6 +29,9 @@ type Snapshot struct { type Dispatcher struct { notifications *executionNotifications + // sessionExecution runs the Session execution operations on the lease; + // Bind sets it from Owner.Sessions. + sessionExecution *sessions.ExecutionOperations Policy Store *store.Store Registry *runtimegateway.Registry diff --git a/services/core/internal/execution/functions.go b/services/core/internal/execution/functions.go index 5f4c4102d..ec278de1a 100644 --- a/services/core/internal/execution/functions.go +++ b/services/core/internal/execution/functions.go @@ -51,6 +51,7 @@ type functionReply struct { type functionExchange struct { store *store.Store + sessions *sessions.ExecutionOperations tenant, session, turn string kind string tools []proto.FunctionTool @@ -70,7 +71,7 @@ func (f *functionExchange) record(ctx context.Context, env proto.Envelope) error if !declared { return errors.New("undeclared function callback") } - err := f.store.RecordFunctionCall(ctx, f.tenant, f.session, f.turn, sessions.FunctionCall{ + err := f.sessions.RecordFunctionCall(ctx, f.tenant, f.session, f.turn, sessions.FunctionCall{ CallID: items.Identity(f.turn, "tool:"+call.CallID), ExecutorCallID: call.CallID, Name: call.Name, Arguments: call.Arguments, }) return f.unlessCancelling(ctx, err) @@ -80,7 +81,7 @@ func (f *functionExchange) start(ctx context.Context, peer *runtimegateway.Sessi if f.reply != nil || len(f.tools) == 0 { return nil } - calls, err := f.store.PendingFunctionCalls(ctx, f.tenant, f.session, f.turn) + calls, err := f.sessions.PendingFunctionCalls(ctx, f.tenant, f.session, f.turn) if err != nil { return err } @@ -119,7 +120,7 @@ func (f *functionExchange) confirm(ctx context.Context, reply functionReply) err if !reply.ack.Applied { return fmt.Errorf("function result not applied: %s", reply.ack.ErrorCode) } - return f.unlessCancelling(ctx, f.store.ConfirmFunctionResult(ctx, f.tenant, f.session, f.turn, id)) + return f.unlessCancelling(ctx, f.sessions.ConfirmFunctionResult(ctx, f.tenant, f.session, f.turn, id)) } func (f *functionExchange) unlessCancelling(ctx context.Context, err error) error { @@ -136,7 +137,7 @@ func (f *functionExchange) complete(ctx context.Context) error { if len(f.tools) == 0 { return nil } - calls, err := f.store.PendingFunctionCalls(ctx, f.tenant, f.session, f.turn) + calls, err := f.sessions.PendingFunctionCalls(ctx, f.tenant, f.session, f.turn) if err != nil { return err } diff --git a/services/core/internal/execution/owner.go b/services/core/internal/execution/owner.go index 38f2bcd2c..e0cf459ed 100644 --- a/services/core/internal/execution/owner.go +++ b/services/core/internal/execution/owner.go @@ -2,9 +2,11 @@ package execution import ( "context" + "errors" "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) @@ -21,6 +23,22 @@ type Owner struct { Lease Ownership Store *store.Store // the remaining store execution operations, built by store.NewExecution(s, lease) Deployment *deployment.ExecutionOperations // sandbox deployment changes on the lease-bound deploymentpg storage + Sessions *sessions.ExecutionOperations // Session execution operations on the lease-bound sessionpg storage +} + +// Bind returns a copy of d bound to owner's execution operations. StartWorker +// binds through it; a Dispatcher that runs Turns outside a Worker binds the same way. +func (d *Dispatcher) Bind(owner Owner) (*Dispatcher, error) { + if owner.Store == nil { + return nil, errors.New("execution requires the execution Store") + } + if owner.Sessions == nil { + return nil, errors.New("execution requires the Session execution operations") + } + owned := *d + owned.Store = owner.Store + owned.sessionExecution = owner.Sessions + return &owned, nil } // leaseCloseTimeout bounds releasing the lease once the Worker owns it. diff --git a/services/core/internal/execution/owner_test.go b/services/core/internal/execution/owner_test.go index 441ff57da..2f5d3eae2 100644 --- a/services/core/internal/execution/owner_test.go +++ b/services/core/internal/execution/owner_test.go @@ -123,19 +123,25 @@ func TestStartWorkerFailureClosesLeaseOnce(t *testing.T) { _, err := StartWorker(canceled, &Dispatcher{Credentials: &recordingCredentials{}, Observer: unusedObserver{t}, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader}, Owner{Lease: lease}) return err }, - "missing deployment": func(t *testing.T, lease *closeCountingLease) error { + "missing Session operations": func(t *testing.T, lease *closeCountingLease) error { _, owner, deployments, deploymentReader := resetManagerStore(t) service, reader := unusedSessions(t) _, err := StartWorker(canceled, &Dispatcher{Credentials: &recordingCredentials{}, Observer: unusedObserver{t}, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader}, Owner{Lease: lease, Store: owner.Store}) return err }, + "missing deployment": func(t *testing.T, lease *closeCountingLease) error { + _, owner, deployments, deploymentReader := resetManagerStore(t) + service, reader := unusedSessions(t) + _, err := StartWorker(canceled, &Dispatcher{Credentials: &recordingCredentials{}, Observer: unusedObserver{t}, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader}, Owner{Lease: lease, Store: owner.Store, Sessions: owner.Sessions}) + return err + }, "deployment claim": func(t *testing.T, lease *closeCountingLease) error { s, owner, deployments, deploymentReader := resetManagerStore(t) lease.inner = owner.Lease id := uuid.NewString() service, reader := unusedSessions(t) dispatcher := &Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), Credentials: &recordingCredentials{}, Observer: unusedObserver{t}, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader, ManagedRuntimes: NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })} - _, err := StartWorker(canceled, dispatcher, Owner{Lease: lease, Store: owner.Store, Deployment: owner.Deployment}) + _, err := StartWorker(canceled, dispatcher, Owner{Lease: lease, Store: owner.Store, Deployment: owner.Deployment, Sessions: owner.Sessions}) if ping := owner.Lease.CheckOwnership(t.Context()); !errors.Is(ping, pgunit.ErrLeaseClosed) { t.Error("failed start kept the database lease", ping) } @@ -159,28 +165,28 @@ func TestStartWorkerChecksDeploymentAfterItsDependencies(t *testing.T) { _, owner, deployments, deploymentReader := resetManagerStore(t) credentials, observer := &recordingCredentials{}, unusedObserver{t} service, reader := unusedSessions(t) + bound := Owner{Store: owner.Store, Sessions: owner.Sessions} for _, test := range []struct { name string dispatcher Dispatcher - store bool + owner Owner want string }{ - {"missing Credentials", Dispatcher{Observer: observer}, true, "execution worker requires MCP Credentials"}, - {"missing observer", Dispatcher{Credentials: credentials}, true, "execution requires a model configuration observer"}, - {"missing deployment service", Dispatcher{Credentials: credentials, Observer: observer}, true, "execution worker requires the deployment service"}, - {"missing deployment reader", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments}, true, "execution worker requires the deployment reader"}, - {"missing Session service", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments, DeploymentReader: deploymentReader}, true, "execution worker requires the Session service"}, - {"missing Session reader", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service}, true, "execution worker requires the Session reader"}, - {"missing Store", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader}, false, "execution worker requires the execution Store"}, - {"missing deployment", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader}, true, "execution worker requires the deployment execution operations"}, + {"missing Credentials", Dispatcher{Observer: observer}, bound, "execution worker requires MCP Credentials"}, + {"missing observer", Dispatcher{Credentials: credentials}, bound, "execution requires a model configuration observer"}, + {"missing deployment service", Dispatcher{Credentials: credentials, Observer: observer}, bound, "execution worker requires the deployment service"}, + {"missing deployment reader", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments}, bound, "execution worker requires the deployment reader"}, + {"missing Session service", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments, DeploymentReader: deploymentReader}, bound, "execution worker requires the Session service"}, + {"missing Session reader", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service}, bound, "execution worker requires the Session reader"}, + {"missing Store", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader}, Owner{Sessions: owner.Sessions}, "execution requires the execution Store"}, + {"missing Session operations", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader}, Owner{Store: owner.Store}, "execution requires the Session execution operations"}, + {"missing deployment", Dispatcher{Credentials: credentials, Observer: observer, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader}, bound, "execution worker requires the deployment execution operations"}, } { t.Run(test.name, func(t *testing.T) { lease := &closeCountingLease{t: t} lease.closable.Store(true) - started := Owner{Lease: lease} - if test.store { - started.Store = owner.Store - } + started := test.owner + started.Lease = lease if _, err := StartWorker(t.Context(), &test.dispatcher, started); err == nil || err.Error() != test.want { t.Fatalf("StartWorker without deployment operations = %v, want %q", err, test.want) } @@ -194,7 +200,7 @@ func TestWorkerRunClosesLeaseAfterDrain(t *testing.T) { id := uuid.NewString() service, reader := unusedSessions(t) dispatcher := &Dispatcher{Store: s, Registry: runtimegateway.NewRegistry(), Credentials: &recordingCredentials{}, Observer: unusedObserver{t}, Deployment: deployments, DeploymentReader: deploymentReader, Sessions: service, SessionsReader: reader, ManagedRuntimes: NewDeferredRuntimeProvider(id, func(context.Context) (*RuntimeProvider, error) { return nil, nil })} - worker, err := StartWorker(t.Context(), dispatcher, Owner{Lease: lease, Store: owner.Store, Deployment: owner.Deployment}) + worker, err := StartWorker(t.Context(), dispatcher, Owner{Lease: lease, Store: owner.Store, Deployment: owner.Deployment, Sessions: owner.Sessions}) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/execution/sandbox_deployment_drain_test.go b/services/core/internal/execution/sandbox_deployment_drain_test.go index 15467107b..5a66a587a 100644 --- a/services/core/internal/execution/sandbox_deployment_drain_test.go +++ b/services/core/internal/execution/sandbox_deployment_drain_test.go @@ -72,7 +72,7 @@ func delayedReadWriter(t *testing.T, armed *atomic.Bool, reading chan struct{}, } }) deployments, reader, operations := testDeployment(t, pool, nil, lease) - return Owner{Lease: lease, Store: store.NewExecution(store.New(pool), lease), Deployment: operations}, deployments, reader, pool + return Owner{Lease: lease, Store: store.NewExecution(store.New(pool), lease), Deployment: operations, Sessions: sessionExecution(t, lease)}, deployments, reader, pool } func TestSandboxDeploymentDrainPreservesLeaseInFlightRead(t *testing.T) { diff --git a/services/core/internal/execution/sandbox_reset_test.go b/services/core/internal/execution/sandbox_reset_test.go index 0c90d6093..7bcc22cfa 100644 --- a/services/core/internal/execution/sandbox_reset_test.go +++ b/services/core/internal/execution/sandbox_reset_test.go @@ -14,11 +14,13 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgtest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimegateway" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/docker" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/e2b" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox/node" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" @@ -65,7 +67,18 @@ func testOwner(t *testing.T, pool *pgxpool.Pool, cipher *credentialcrypto.Cipher } t.Cleanup(func() { _ = lease.Close(context.Background()) }) deployments, reader, operations := testDeployment(t, pool, cipher, lease) - return Owner{Lease: lease, Store: store.NewExecution(s, lease), Deployment: operations}, deployments, reader + return Owner{Lease: lease, Store: store.NewExecution(s, lease), Deployment: operations, Sessions: sessionExecution(t, lease)}, deployments, reader +} + +// sessionExecution builds the Session execution operations on lease, as +// cmd/server does. +func sessionExecution(t *testing.T, lease *pgunit.Lease) *sessions.ExecutionOperations { + t.Helper() + operations, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) + if err != nil { + t.Fatal(err) + } + return operations } func TestSandboxResetPageTimeoutRecoversCommittedOwner(t *testing.T) { diff --git a/services/core/internal/execution/worker.go b/services/core/internal/execution/worker.go index b7ccc1f88..f42cd8589 100644 --- a/services/core/internal/execution/worker.go +++ b/services/core/internal/execution/worker.go @@ -68,16 +68,15 @@ func StartWorker(ctx context.Context, dispatcher *Dispatcher, owner Owner) (_ *W if dispatcher.SessionsReader == nil { return nil, errors.New("execution worker requires the Session reader") } - if owner.Store == nil { - return nil, errors.New("execution worker requires the execution Store") + owned, err := dispatcher.Bind(owner) + if err != nil { + return nil, err } if owner.Deployment == nil { return nil, errors.New("execution worker requires the deployment execution operations") } - owned := *dispatcher - owned.Store = owner.Store owned.notifications = &executionNotifications{} - worker := &Worker{concurrency: dispatcher.MaxConcurrentExecutions, dispatcher: &owned, admission: dispatcher.Store, lease: owner.Lease, directoryReads: make(chan directoryReadRequest), fileWrites: make(chan fileWriteRequest), stopped: make(chan struct{}), scheduleWake: make(chan struct{}, 1), enrolledConnections: make(map[string]*runtimeConnection)} + worker := &Worker{concurrency: dispatcher.MaxConcurrentExecutions, dispatcher: owned, admission: dispatcher.Store, lease: owner.Lease, directoryReads: make(chan directoryReadRequest), fileWrites: make(chan fileWriteRequest), stopped: make(chan struct{}), scheduleWake: make(chan struct{}, 1), enrolledConnections: make(map[string]*runtimeConnection)} worker.runtimes, err = newRuntimeManager(owner, owned.Deployment, owned.DeploymentReader, owned.Registry, owned.ManagedRuntimes) if err != nil { return nil, err diff --git a/services/core/internal/persistence/postgres/sessionpg/execution.go b/services/core/internal/persistence/postgres/sessionpg/execution.go new file mode 100644 index 000000000..ed2246f4c --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution.go @@ -0,0 +1,17 @@ +package sessionpg + +import ( + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// Execution is the Session storage of the execution owner. Every transaction +// runs on the connection that holds the execution lease, which the owner +// acquired and closes; losing or closing the lease fails the operation and +// never falls back to the pool. +type Execution struct{ lease *pgunit.Lease } + +// NewExecution builds the Session execution storage on a lease it borrows. +func NewExecution(lease *pgunit.Lease) *Execution { return &Execution{lease: lease} } + +var _ sessions.ExecutionStorage = (*Execution)(nil) diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_functions.go b/services/core/internal/persistence/postgres/sessionpg/execution_functions.go new file mode 100644 index 000000000..2f539061c --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution_functions.go @@ -0,0 +1,28 @@ +package sessionpg + +import ( + "context" + "errors" + + "github.com/jackc/pgx/v5" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +func (e *Execution) WithFunctionTurn(ctx context.Context, tenant, session, turn string, apply func(context.Context, sessions.FunctionTx, sessions.Turn) error) error { + p, err := TurnLookup(tenant, session, turn) + if err != nil { + return err + } + return WithSession(ctx, e.lease, p.TenantID, p.SessionID, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { + row, err := q.GetTurn(ctx, p) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.ErrNotFound + } + if err != nil { + return err + } + return apply(ctx, BindSession(q, p.TenantID, p.SessionID), TurnFromRow(row)) + }) +} diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_functions_test.go b/services/core/internal/persistence/postgres/sessionpg/execution_functions_test.go new file mode 100644 index 000000000..435cde321 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/execution_functions_test.go @@ -0,0 +1,266 @@ +package sessionpg + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "reflect" + "slices" + "strings" + "testing" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" + "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" +) + +// sessionExecution takes the execution lease of pool's database and builds +// the Session execution operations on it. The lease closes when the test ends. +func sessionExecution(t *testing.T, pool *pgxpool.Pool) (*sessions.ExecutionOperations, *pgunit.Lease) { + t.Helper() + lease, err := pgunit.AcquireLease(t.Context(), pool) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = lease.Close(context.Background()) }) + operations, err := sessions.NewExecutionOperations(NewExecution(lease)) + if err != nil { + t.Fatal(err) + } + return operations, lease +} + +func functionCall(id string) sessions.FunctionCall { + return sessions.FunctionCall{CallID: id, ExecutorCallID: "native-" + id, Name: "lookup", Arguments: json.RawMessage(`{"ticket":9007199254740993}`)} +} + +func turnStatus(t *testing.T, pool *pgxpool.Pool, f functionTurn) string { + t.Helper() + var status string + if err := pool.QueryRow(t.Context(), `SELECT status FROM turns WHERE id = $1`, f.turn).Scan(&status); err != nil { + t.Fatal(err) + } + return status +} + +// Calls, results and receipts commit on the lease, survive a new execution +// owner, and journal the Turn's required actions after each change. +func TestFunctionExecutionRecordsCallsAndReceipts(t *testing.T) { + pool := pgtest.OpenIsolated(t, nil) + operations, lease := sessionExecution(t, pool) + f := newFunctionTurn(t, pool, sessions.TurnInProgress) + tenant, session, turn := f.ids() + results := []string{ + `{"success":true,"output":"answer"}`, + `{"success":true,"output":""}`, + `{"success":false,"output":[{"type":"input_text","text":"before"},{"type":"input_image","image_url":"data:image/png;base64,AA=="},{"type":"input_text","text":""}],"error":"tool failed"}`, + `{"success":true,"output":[],"error":null}`, + `{"success":false,"output":null,"error":"missing"}`, + `{"success":false}`, + } + for i, raw := range results { + call := functionCall(fmt.Sprint(i)) + if err := operations.RecordFunctionCall(t.Context(), tenant, session, turn, call); err != nil { + t.Fatal(err) + } + retry := call + retry.Arguments = json.RawMessage(`{ "ticket" : 9007199254740993 }`) + if err := operations.RecordFunctionCall(t.Context(), tenant, session, turn, retry); err != nil { + t.Fatal("reformatted retry", err) + } + if err := operations.ConfirmFunctionResult(t.Context(), tenant, session, turn, call.CallID); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatal("unsubmitted result applied", err) + } + if _, err := admitResult(t.Context(), pool, f, turn, call.CallID, raw); err != nil { + t.Fatal(err) + } + } + if status := turnStatus(t, pool, f); status != sessions.TurnWaiting { + t.Fatal(status) + } + before, err := operations.PendingFunctionCalls(t.Context(), tenant, session, turn) + if err != nil || len(before) != len(results) { + t.Fatal(before, err) + } + + // A new execution owner reads the same calls and results. + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, pool) + if err := lease.Close(t.Context()); err != nil { + t.Fatal(err) + } + awaitRelease() + if err := operations.RecordFunctionCall(t.Context(), tenant, session, turn, functionCall("late")); !errors.Is(err, pgunit.ErrLeaseClosed) { + t.Fatal("closed lease recorded a call", err) + } + operations, _ = sessionExecution(t, pool) + restored, err := operations.PendingFunctionCalls(t.Context(), tenant, session, turn) + if err != nil || !reflect.DeepEqual(restored, before) { + t.Fatal("recovery lost calls or results", restored, err) + } + for i, expected := range results { + id := fmt.Sprint(i) + call := restored[i] + if call.CallID != id || call.Applied || call.ExecutorCallID != "native-"+id || !strings.Contains(string(call.Arguments), "9007199254740993") { + t.Fatal(call) + } + saved, err := jsonobject.Normalize(call.Result) + wanted, _ := jsonobject.Normalize(json.RawMessage(expected)) + if err != nil || string(saved) != string(wanted) { + t.Fatal("result changed", string(call.Result), err) + } + if _, err := admitResult(t.Context(), pool, f, turn, id, expected); err != nil { + t.Fatal(err) + } + if _, err := admitResult(t.Context(), pool, f, turn, id, `{"success":false,"error":"changed"}`); !errors.Is(err, sessions.ErrFunctionResultConflict) { + t.Fatal(err) + } + for range 2 { + if err := operations.ConfirmFunctionResult(t.Context(), tenant, session, turn, id); err != nil { + t.Fatal(err) + } + } + } + pending, err := operations.PendingFunctionCalls(t.Context(), tenant, session, turn) + if err != nil || pending == nil || len(pending) != 0 { + t.Fatal(pending, err) + } + if status := turnStatus(t, pool, f); status != sessions.TurnInProgress { + t.Fatal(status) + } + if row := storedCall(t, pool, f, f.turn, "0"); !row.Applied { + t.Fatal("receipt did not persist") + } + + // Each change journals the required actions once: six calls, then five + // receipts that leave actions, then the resumed Turn and Session. + _, changes := journal(t, pool, f.session) + want := append(slices.Repeat([]string{"agent.session.requires_action"}, 11), "agent.session.turn.in_progress", "agent.session.in_progress") + if got := kinds(changes); !slices.Equal(got, want) { + t.Fatalf("journal %v", got) + } + for i, change := range changes[:11] { + count := i + 1 + if i >= 6 { + count = 11 - i + } + if len(change.RequiredActions) != count || change.Turn == nil || change.Turn.Status != sessions.TurnWaiting { + t.Fatalf("change %d: %+v", i, change) + } + } + var precise int + if err := pool.QueryRow(t.Context(), `SELECT count(*) FROM session_events WHERE session_id = $1 AND payload::text LIKE '%"ticket": 9007199254740993%'`, f.session).Scan(&precise); err != nil || precise != 11 { + t.Fatal("argument precision lost", precise, err) + } +} + +// Function calls are scoped to the tenant's Turn, keep their identities, and +// change only while the Turn accepts results. +func TestFunctionExecutionScopesAndRejects(t *testing.T) { + pool := pgtest.OpenIsolated(t, nil) + operations, _ := sessionExecution(t, pool) + + t.Run("identities are immutable", func(t *testing.T) { + f := newFunctionTurn(t, pool, sessions.TurnQueued) + tenant, session, turn := f.ids() + call := functionCall("call") + if err := operations.RecordFunctionCall(t.Context(), tenant, session, turn, call); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatal("queued Turn accepted a call", err) + } + exec(t, pool, `UPDATE turns SET status = 'in_progress' WHERE id = $1`, f.turn) + if err := operations.RecordFunctionCall(t.Context(), tenant, session, turn, call); err != nil { + t.Fatal(err) + } + for field, change := range map[string]func(*sessions.FunctionCall){ + "call": func(c *sessions.FunctionCall) { c.CallID = "another-public-id" }, + "executor": func(c *sessions.FunctionCall) { c.ExecutorCallID = "another-executor-id" }, + "name": func(c *sessions.FunctionCall) { c.Name = "another-function" }, + "arguments": func(c *sessions.FunctionCall) { c.Arguments = json.RawMessage(`{"ticket":9007199254740992}`) }, + } { + changed := call + change(&changed) + if err := operations.RecordFunctionCall(t.Context(), tenant, session, turn, changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { + t.Fatal(field, err) + } + } + // The same public identity is a new call in another Turn. + exec(t, pool, `UPDATE turns SET status = 'completed', completed_at = clock_timestamp() WHERE id = $1`, f.turn) + next := f + next.turn = addTurn(t, pool, f.session, sessions.TurnInProgress) + if err := operations.RecordFunctionCall(t.Context(), tenant, session, next.turnID(), call); err != nil { + t.Fatal("call identity leaked across Turns", err) + } + }) + + t.Run("lookups are scoped", func(t *testing.T) { + f := newFunctionTurn(t, pool, sessions.TurnInProgress) + tenant, session, turn := f.ids() + if err := operations.RecordFunctionCall(t.Context(), tenant, session, turn, functionCall("call")); err != nil { + t.Fatal(err) + } + for _, scope := range []struct{ tenant, session, turn string }{ + {uuid.NewString(), session, turn}, {tenant, uuid.NewString(), turn}, {tenant, session, uuid.NewString()}, + } { + if _, err := operations.PendingFunctionCalls(t.Context(), scope.tenant, scope.session, scope.turn); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal(err) + } + if err := operations.RecordFunctionCall(t.Context(), scope.tenant, scope.session, scope.turn, functionCall("call")); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal(err) + } + if err := operations.ConfirmFunctionResult(t.Context(), scope.tenant, scope.session, scope.turn, "call"); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal(err) + } + } + for _, scope := range []struct{ tenant, session, turn string }{ + {"tenant", session, turn}, {tenant, "session", turn}, {tenant, session, "turn"}, + } { + if _, err := operations.PendingFunctionCalls(t.Context(), scope.tenant, scope.session, scope.turn); !errors.Is(err, sessions.ErrInvalidInput) { + t.Fatal(err) + } + } + if err := operations.ConfirmFunctionResult(t.Context(), tenant, session, turn, "missing"); !errors.Is(err, sessions.ErrNotFound) { + t.Fatal(err) + } + // A publicly deleted Session still settles its Turn's calls. + exec(t, pool, `UPDATE sessions SET deleted_at = clock_timestamp() WHERE id = $1`, f.session) + if pending, err := operations.PendingFunctionCalls(t.Context(), tenant, session, turn); err != nil || len(pending) != 1 { + t.Fatal(pending, err) + } + }) + + for _, end := range []string{"cancelling", sessions.TurnCancelled, sessions.TurnCompleted, sessions.TurnFailed} { + t.Run("no changes after "+end, func(t *testing.T) { + f := newFunctionTurn(t, pool, sessions.TurnInProgress) + tenant, session, turn := f.ids() + for _, id := range []string{"submitted", "pending"} { + if err := operations.RecordFunctionCall(t.Context(), tenant, session, turn, functionCall(id)); err != nil { + t.Fatal(err) + } + } + if _, err := admitResult(t.Context(), pool, f, turn, "submitted", `{"success":true,"output":"saved"}`); err != nil { + t.Fatal(err) + } + if end == "cancelling" { + exec(t, pool, `UPDATE turns SET cancel_requested_at = clock_timestamp() WHERE id = $1`, f.turn) + } else { + exec(t, pool, `UPDATE turns SET status = $2, completed_at = clock_timestamp() WHERE id = $1`, f.turn, end) + } + if pending, err := operations.PendingFunctionCalls(t.Context(), tenant, session, turn); err != nil || len(pending) != 0 { + t.Fatal(pending, err) + } + if err := operations.ConfirmFunctionResult(t.Context(), tenant, session, turn, "submitted"); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatal(err) + } + if err := operations.RecordFunctionCall(t.Context(), tenant, session, turn, functionCall("late")); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatal(err) + } + if row := storedCall(t, pool, f, f.turn, "submitted"); row.Applied || len(row.Result) == 0 { + t.Fatal("saved result lost or falsely applied") + } + }) + } +} diff --git a/services/core/internal/persistence/postgres/sessionpg/functions.go b/services/core/internal/persistence/postgres/sessionpg/functions.go new file mode 100644 index 000000000..8da77be69 --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/functions.go @@ -0,0 +1,167 @@ +package sessionpg + +import ( + "context" + "encoding/json" + "errors" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +var ( + _ sessions.FunctionResultTx = (*SessionTx)(nil) + _ sessions.FunctionTx = (*SessionTx)(nil) +) + +// LoadRequiredActions reads the required actions of the Session's Turn within +// the caller's snapshot, as sessions.LoadRequiredActions selects them. +func LoadRequiredActions(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, turn sessions.Turn) ([]v1.FunctionCallAction, error) { + return sessions.LoadRequiredActions(ctx, pendingCalls{q: q, session: session}, turn) +} + +// pendingCalls reads a Session's pending function calls without the Session +// lock, for a snapshot read. +type pendingCalls struct { + q *sqlc.Queries + session pgtype.UUID +} + +func (p pendingCalls) LoadPendingFunctionCalls(ctx context.Context, turn string) ([]sessions.FunctionCall, error) { + return loadPendingFunctionCalls(ctx, p.q, p.session, turn) +} + +func loadPendingFunctionCalls(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, turn string) ([]sessions.FunctionCall, error) { + id, err := parseID(turn) + if err != nil { + return nil, err + } + rows, err := q.ListPendingFunctionCalls(ctx, sqlc.ListPendingFunctionCallsParams{SessionID: session, TurnID: id}) + if err != nil { + return nil, err + } + calls := make([]sessions.FunctionCall, 0, len(rows)) + for _, row := range rows { + calls = append(calls, functionCallFromRow(row)) + } + return calls, nil +} + +func functionCallFromRow(row sqlc.FunctionCall) sessions.FunctionCall { + return sessions.FunctionCall{CallID: row.CallID, ExecutorCallID: row.ExecutorCallID, Name: row.Name, Arguments: row.Arguments, Result: row.Result, Applied: row.Applied} +} + +func (t *SessionTx) LoadPendingFunctionCalls(ctx context.Context, turn string) ([]sessions.FunctionCall, error) { + return loadPendingFunctionCalls(ctx, t.q, t.session, turn) +} + +func (t *SessionTx) FindResultTurn(ctx context.Context, turn string) (sessions.Turn, bool, error) { + row, err := t.q.SessionEventTurn(ctx, sqlc.SessionEventTurnParams{SessionID: t.session, ID: pgunit.PathID(turn)}) + 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) MatchFunctionResult(ctx context.Context, turn, call string, result json.RawMessage) (sessions.FunctionResultMatch, error) { + id, err := parseID(turn) + if err != nil { + return sessions.FunctionResultMatch{}, err + } + row, err := t.q.MatchFunctionResult(ctx, sqlc.MatchFunctionResultParams{SessionID: t.session, TurnID: id, CallID: call, Result: result}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.FunctionResultMatch{}, nil + } + if err != nil { + return sessions.FunctionResultMatch{}, err + } + return sessions.FunctionResultMatch{Recorded: true, Submitted: row.Submitted, Matches: row.Matches}, nil +} + +func (t *SessionTx) SubmitFunctionResult(ctx context.Context, turn, call string, result json.RawMessage) error { + id, err := parseID(turn) + if err != nil { + return err + } + return t.q.SubmitFunctionResult(ctx, sqlc.SubmitFunctionResultParams{SessionID: t.session, TurnID: id, CallID: call, Result: result}) +} + +func (t *SessionTx) HasFunctionCall(ctx context.Context, call string) (bool, error) { + return t.q.SessionHasFunctionCall(ctx, sqlc.SessionHasFunctionCallParams{SessionID: t.session, CallID: call}) +} + +func (t *SessionTx) MatchFunctionCall(ctx context.Context, turn string, call sessions.FunctionCall) (sessions.FunctionCallMatch, error) { + id, err := parseID(turn) + if err != nil { + return sessions.FunctionCallMatch{}, err + } + matches, err := t.q.MatchFunctionCall(ctx, sqlc.MatchFunctionCallParams{SessionID: t.session, TurnID: id, CallID: call.CallID, ExecutorCallID: call.ExecutorCallID, Name: call.Name, Arguments: call.Arguments}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.FunctionCallMatch{}, nil + } + if err != nil { + return sessions.FunctionCallMatch{}, err + } + return sessions.FunctionCallMatch{Recorded: true, Matches: matches}, nil +} + +func (t *SessionTx) CreateFunctionCall(ctx context.Context, turn string, call sessions.FunctionCall) error { + id, err := parseID(turn) + if err != nil { + return err + } + count, err := t.q.CreateFunctionCall(ctx, sqlc.CreateFunctionCallParams{SessionID: t.session, TurnID: id, CallID: call.CallID, ExecutorCallID: call.ExecutorCallID, Name: call.Name, Arguments: call.Arguments}) + if err != nil { + return err + } + if count != 1 { + return sessions.ErrIdempotencyConflict + } + return nil +} + +func (t *SessionTx) LoadFunctionCall(ctx context.Context, turn, call string) (sessions.FunctionCall, bool, error) { + id, err := parseID(turn) + if err != nil { + return sessions.FunctionCall{}, false, err + } + row, err := t.q.GetFunctionCall(ctx, sqlc.GetFunctionCallParams{TenantID: t.tenant, SessionID: t.session, TurnID: id, CallID: call}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.FunctionCall{}, false, nil + } + if err != nil { + return sessions.FunctionCall{}, false, err + } + return functionCallFromRow(row), true, nil +} + +func (t *SessionTx) ApplyFunctionResult(ctx context.Context, turn, call string) error { + id, err := parseID(turn) + if err != nil { + return err + } + return t.q.ApplyFunctionResult(ctx, sqlc.ApplyFunctionResultParams{SessionID: t.session, TurnID: id, CallID: call}) +} + +func (t *SessionTx) ApplyFunctionTurnStatus(ctx context.Context, turn, expected, status string) (sessions.Turn, error) { + id, err := parseID(turn) + if err != nil { + return sessions.Turn{}, err + } + row, err := t.q.TransitionTurn(ctx, sqlc.TransitionTurnParams{ID: id, SessionID: t.session, ExpectedStatus: expected, NewStatus: status, Outcome: []byte(`{}`)}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.Turn{}, sessions.ErrTurnConflict + } + if err != nil { + return sessions.Turn{}, err + } + return TurnFromRow(row), nil +} diff --git a/services/core/internal/persistence/postgres/sessionpg/functions_test.go b/services/core/internal/persistence/postgres/sessionpg/functions_test.go new file mode 100644 index 000000000..66c2aef1a --- /dev/null +++ b/services/core/internal/persistence/postgres/sessionpg/functions_test.go @@ -0,0 +1,214 @@ +package sessionpg + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "testing" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/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" +) + +// functionTurn is a stored tenant's Session with one Turn. +type functionTurn struct { + tenant, session, turn pgtype.UUID +} + +func (f functionTurn) ids() (string, string, string) { + return uuid.UUID(f.tenant.Bytes).String(), uuid.UUID(f.session.Bytes).String(), uuid.UUID(f.turn.Bytes).String() +} + +func (f functionTurn) turnID() string { return uuid.UUID(f.turn.Bytes).String() } + +// newFunctionTurn stores a fresh tenant's Session with one Turn of status. +func newFunctionTurn(t *testing.T, pool *pgxpool.Pool, status string) functionTurn { + t.Helper() + tenant, session := uuid.New(), uuid.New() + exec(t, pool, `INSERT INTO sessions(id, tenant_id, engine, idempotency_key, request_hash) VALUES ($1, $2, 'codex', 'key', 'hash')`, session, tenant) + f := functionTurn{tenant: pgID(tenant), session: pgID(session)} + f.turn = addTurn(t, pool, f.session, status) + return f +} + +// addTurn stores another Turn of status in the Session. +func addTurn(t *testing.T, pool *pgxpool.Pool, session pgtype.UUID, status string) pgtype.UUID { + t.Helper() + turn := uuid.New() + exec(t, pool, `INSERT INTO turns(id, session_id, status, completed_at) + VALUES ($1, $2, $3::text, CASE WHEN $3::text IN ('completed', 'failed', 'cancelled') THEN clock_timestamp() END)`, turn, session, status) + return pgID(turn) +} + +// storeCall stores a function call of the Turn, with result when it is not +// empty. +func storeCall(t *testing.T, pool *pgxpool.Pool, session, turn pgtype.UUID, call, result string) { + t.Helper() + var saved any + if result != "" { + saved = json.RawMessage(result) + } + exec(t, pool, `INSERT INTO function_calls(session_id, turn_id, call_id, executor_call_id, name, arguments, result) VALUES ($1, $2, $3, 'native-' || $3, 'lookup', '{}', $4::jsonb)`, session, turn, call, saved) +} + +// storedCall reads a stored function call. +func storedCall(t *testing.T, pool *pgxpool.Pool, f functionTurn, turn pgtype.UUID, call string) sqlc.FunctionCall { + t.Helper() + row, err := sqlc.New(pool).GetFunctionCall(t.Context(), sqlc.GetFunctionCallParams{TenantID: f.tenant, SessionID: f.session, TurnID: turn, CallID: call}) + if err != nil { + t.Fatal(err) + } + return row +} + +// admitResult admits a function result in its own pooled Session transaction, +// as input admission does. +func admitResult(ctx context.Context, pool *pgxpool.Pool, f functionTurn, turn, call, result string) (sessions.Turn, error) { + var admitted sessions.Turn + err := WithSession(ctx, pgunit.NewPool(pool), f.tenant, f.session, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { + var err error + admitted, err = sessions.AdmitFunctionResult(ctx, BindSession(q, f.tenant, f.session), sessions.FunctionResultInput{TurnID: turn, CallID: call, Result: json.RawMessage(result)}) + return err + }) + return admitted, err +} + +func TestAdmitFunctionResultThroughTheBinding(t *testing.T) { + pool := pgtest.Open(t) + f := newFunctionTurn(t, pool, sessions.TurnWaiting) + storeCall(t, pool, f.session, f.turn, "call", "") + other := addTurn(t, pool, f.session, sessions.TurnCompleted) + storeCall(t, pool, f.session, other, "earlier", `{"success":true}`) + storeCall(t, pool, f.session, other, "late", "") + // Another Session's calls never classify this Session's results. + foreign := newFunctionTurn(t, pool, sessions.TurnWaiting) + storeCall(t, pool, foreign.session, foreign.turn, "foreign", "") + + turn, err := admitResult(t.Context(), pool, f, f.turnID(), "call", `{"success":true,"output":"saved"}`) + if err != nil || turn.ID != f.turnID() || turn.Status != sessions.TurnWaiting { + t.Fatal(turn, err) + } + if row := storedCall(t, pool, f, f.turn, "call"); string(row.Result) != `{"output": "saved", "success": true}` || row.Applied { + t.Fatalf("stored %s", row.Result) + } + if _, err := admitResult(t.Context(), pool, f, f.turnID(), "call", `{"output":"saved","success":true}`); err != nil { + t.Fatal("identical retry", err) + } + if _, err := admitResult(t.Context(), pool, f, f.turnID(), "call", `{"success":false}`); !errors.Is(err, sessions.ErrFunctionResultConflict) { + t.Fatal("changed result", err) + } + // A result saved before the Turn ended still answers an identical retry. + if _, err := admitResult(t.Context(), pool, f, uuid.UUID(other.Bytes).String(), "earlier", `{"success":true}`); err != nil { + t.Fatal("retry after the Turn ended", err) + } + for name, test := range map[string]struct { + turn, call string + want error + }{ + "call of another Turn": {f.turnID(), "earlier", sessions.ErrFunctionCallTurnMismatch}, + "unknown call": {f.turnID(), "missing", sessions.ErrUnknownFunctionCall}, + "unknown Turn": {uuid.NewString(), "call", sessions.ErrFunctionCallTurnMismatch}, + "malformed Turn": {"not-a-turn", "call", sessions.ErrFunctionCallTurnMismatch}, + "malformed Turn and unknown call": {"not-a-turn", "missing", sessions.ErrUnknownFunctionCall}, + "another Session's call": {f.turnID(), "foreign", sessions.ErrUnknownFunctionCall}, + "another Session's Turn and call": {foreign.turnID(), "foreign", sessions.ErrUnknownFunctionCall}, + "another Session's Turn, a call here": {foreign.turnID(), "call", sessions.ErrFunctionCallTurnMismatch}, + "new result for an ended Turn": {uuid.UUID(other.Bytes).String(), "late", sessions.ErrTurnConflict}, + } { + if _, err := admitResult(t.Context(), pool, f, test.turn, test.call, `{"success":true}`); !errors.Is(err, test.want) { + t.Errorf("%s: %v", name, err) + } + } + + // A Turn whose cancellation was requested takes no new result. + storeCall(t, pool, f.session, f.turn, "pending", "") + exec(t, pool, `UPDATE turns SET cancel_requested_at = clock_timestamp() WHERE id = $1`, f.turn) + if _, err := admitResult(t.Context(), pool, f, f.turnID(), "pending", `{"success":true}`); !errors.Is(err, sessions.ErrTurnConflict) { + t.Fatal("cancelling Turn", err) + } + if row := storedCall(t, pool, f, f.turn, "pending"); row.Result != nil { + t.Fatalf("stored %s", row.Result) + } +} + +func TestAdmitFunctionResultRollsBackWithItsTransaction(t *testing.T) { + pool := pgtest.Open(t) + f := newFunctionTurn(t, pool, sessions.TurnWaiting) + storeCall(t, pool, f.session, f.turn, "call", "") + err := WithSession(t.Context(), pgunit.NewPool(pool), f.tenant, f.session, func(ctx context.Context, q *sqlc.Queries, _ sessions.LockedSession) error { + if _, err := sessions.AdmitFunctionResult(ctx, BindSession(q, f.tenant, f.session), sessions.FunctionResultInput{TurnID: f.turnID(), CallID: "call", Result: json.RawMessage(`{"success":true}`)}); err != nil { + return err + } + return errRollback + }) + if row := storedCall(t, pool, f, f.turn, "call"); !errors.Is(err, errRollback) || row.Result != nil { + t.Fatalf("rolled back result kept: %s %v", row.Result, err) + } +} + +func TestConcurrentFunctionResultsChooseOneValue(t *testing.T) { + pool := pgtest.Open(t) + f := newFunctionTurn(t, pool, sessions.TurnWaiting) + storeCall(t, pool, f.session, f.turn, "call", "") + var wg sync.WaitGroup + outcomes := make(chan error, 2) + for _, output := range []string{"first", "second"} { + wg.Go(func() { + _, err := admitResult(t.Context(), pool, f, f.turnID(), "call", fmt.Sprintf(`{"success":true,"output":%q}`, output)) + outcomes <- err + }) + } + wg.Wait() + close(outcomes) + winners, conflicts := 0, 0 + for err := range outcomes { + switch { + case err == nil: + winners++ + case errors.Is(err, sessions.ErrFunctionResultConflict): + conflicts++ + default: + t.Fatal(err) + } + } + if winners != 1 || conflicts != 1 { + t.Fatal(winners, conflicts) + } +} + +func TestLoadRequiredActionsReadsPendingCalls(t *testing.T) { + pool := pgtest.Open(t) + q := sqlc.New(pool) + f := newFunctionTurn(t, pool, sessions.TurnWaiting) + for _, call := range []string{"b", "a", "applied"} { + storeCall(t, pool, f.session, f.turn, call, "") + } + exec(t, pool, `UPDATE function_calls SET created_at = '2026-01-01', arguments = '{"ticket":9007199254740993}' WHERE session_id = $1 AND call_id IN ('a', 'b')`, f.session) + exec(t, pool, `UPDATE function_calls SET result = '{}', applied = true WHERE session_id = $1 AND call_id = 'applied'`, f.session) + row, err := q.GetTurn(t.Context(), sqlc.GetTurnParams{TenantID: f.tenant, SessionID: f.session, ID: f.turn}) + if err != nil { + t.Fatal(err) + } + actions, err := LoadRequiredActions(t.Context(), q, f.session, TurnFromRow(row)) + if err != nil || len(actions) != 2 || actions[0].CallID != "a" || actions[1].CallID != "b" { + t.Fatalf("actions %+v %v", actions, err) + } + encoded, err := json.Marshal(actions[0]) + if err != nil || string(encoded) != `{"arguments":{"ticket":9007199254740993},"call_id":"a","name":"lookup","turn_id":"`+f.turnID()+`","type":"function_call"}` { + t.Fatal(string(encoded), err) + } + exec(t, pool, `UPDATE turns SET cancel_requested_at = clock_timestamp() WHERE id = $1`, f.turn) + // The Turn as the caller read it still waits; the read of its calls sees + // the requested cancellation. + if actions, err := LoadRequiredActions(t.Context(), q, f.session, TurnFromRow(row)); err != nil || actions == nil || len(actions) != 0 { + t.Fatalf("cancelling Turn actions %+v %v", actions, err) + } +} diff --git a/services/core/internal/sessions/doc.go b/services/core/internal/sessions/doc.go index 690c0ac9a..df8e82db8 100644 --- a/services/core/internal/sessions/doc.go +++ b/services/core/internal/sessions/doc.go @@ -6,7 +6,11 @@ // 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. Service runs the pooled Session use cases, such as staging a Turn's +// and terminating an Environment, tracking input activity, the admission +// gates, admitting a function result and reading a Turn's required actions. +// Service runs the pooled Session use cases, such as staging a Turn's // Artifacts, over Storage, and Reader declares the Session reads. +// ExecutionOperations runs the Session writes only the execution owner makes, +// such as recording function calls and their application receipts, over the +// lease-bound ExecutionStorage. package sessions diff --git a/services/core/internal/sessions/execution.go b/services/core/internal/sessions/execution.go new file mode 100644 index 000000000..8a3a8869c --- /dev/null +++ b/services/core/internal/sessions/execution.go @@ -0,0 +1,22 @@ +package sessions + +import "errors" + +// ExecutionOperations are the Session writes only the execution owner makes. +// Each runs in a Session transaction on the connection that holds the +// execution lease, so losing the lease fences it. +type ExecutionOperations struct{ storage ExecutionStorage } + +// NewExecutionOperations binds the Session execution operations to the +// lease-bound storage of one execution owner. +func NewExecutionOperations(storage ExecutionStorage) (*ExecutionOperations, error) { + if storage == nil { + return nil, errors.New("session execution operations require execution storage") + } + return &ExecutionOperations{storage: storage}, nil +} + +// ExecutionStorage is the lease-bound Session storage, one family per line. +type ExecutionStorage interface { + FunctionExecution +} diff --git a/services/core/internal/sessions/execution_functions.go b/services/core/internal/sessions/execution_functions.go new file mode 100644 index 000000000..aa6478c14 --- /dev/null +++ b/services/core/internal/sessions/execution_functions.go @@ -0,0 +1,165 @@ +package sessions + +import "context" + +// FunctionCallMatch is what a Turn's stored call shows about a reported call. +type FunctionCallMatch struct { + // Recorded reports that the Turn already has a call with the public + // identity. + Recorded bool + // Matches reports that the recorded call has the same executor identity, + // name and arguments. + Matches bool +} + +// FunctionTx is one Turn's function calls inside a Session transaction on the +// execution lease. +type FunctionTx interface { + JournalTx + PendingFunctionCallsTx + // MatchFunctionCall compares call with the Turn's call of the same public + // identity. + MatchFunctionCall(ctx context.Context, turn string, call FunctionCall) (FunctionCallMatch, error) + // CreateFunctionCall records a new call of the Turn. A call the Turn + // already has is ErrIdempotencyConflict. + CreateFunctionCall(ctx context.Context, turn string, call FunctionCall) error + // LoadFunctionCall reads the Turn's call and reports whether it has it. + LoadFunctionCall(ctx context.Context, turn, call string) (FunctionCall, bool, error) + // ApplyFunctionResult records the application receipt of the call's + // result. + ApplyFunctionResult(ctx context.Context, turn, call string) error + // ApplyFunctionTurnStatus moves the Turn from expected to status and + // returns it as stored. A Turn no longer in expected is ErrTurnConflict. + ApplyFunctionTurnStatus(ctx context.Context, turn, expected, status string) (Turn, error) +} + +// FunctionExecution is the lease-bound storage of function calls. +type FunctionExecution interface { + // WithFunctionTurn runs apply in one transaction on the execution lease, + // under the Session lock, with the tenant's Turn as the transaction reads + // it. It commits only when apply succeeds. A malformed ID is + // ErrInvalidInput, and a missing Session or Turn is ErrNotFound; a publicly + // deleted Session still settles its Turn. + WithFunctionTurn(ctx context.Context, tenant, session, turn string, apply func(context.Context, FunctionTx, Turn) error) error +} + +// decideFunctionCall decides whether to record a call an executor reported: +// a retry of the recorded call records nothing, a different call with the same +// public identity conflicts, and a new call needs a Turn that accepts results. +func decideFunctionCall(turn Turn, match FunctionCallMatch) (bool, error) { + if match.Recorded { + if !match.Matches { + return false, ErrIdempotencyConflict + } + return false, nil + } + if !acceptsFunctionResults(turn) { + return false, ErrTurnConflict + } + return true, nil +} + +// decideFunctionConfirmation decides whether to record the native application +// of a call's result: an applied call records nothing, and a call without a +// result or on a Turn that no longer accepts results conflicts. +func decideFunctionConfirmation(turn Turn, call FunctionCall) (bool, error) { + if call.Applied { + return false, nil + } + if len(call.Result) == 0 || !acceptsFunctionResults(turn) { + return false, ErrTurnConflict + } + return true, nil +} + +// RecordFunctionCall records a call an executor reported for its Turn and +// journals the Turn's required actions with it. A retry of the recorded call +// changes nothing. +func (o *ExecutionOperations) RecordFunctionCall(ctx context.Context, tenant, session, turn string, call FunctionCall) error { + if err := validateFunctionCall(call); err != nil { + return err + } + return o.storage.WithFunctionTurn(ctx, tenant, session, turn, func(ctx context.Context, tx FunctionTx, current Turn) error { + match, err := tx.MatchFunctionCall(ctx, current.ID, call) + if err != nil { + return err + } + record, err := decideFunctionCall(current, match) + if err != nil || !record { + return err + } + if err := tx.CreateFunctionCall(ctx, current.ID, call); err != nil { + return err + } + return settleFunctionState(ctx, tx, current) + }) +} + +// PendingFunctionCalls returns the Turn's calls without an application +// receipt, with their saved results. A cancelling or ended Turn has none. +func (o *ExecutionOperations) PendingFunctionCalls(ctx context.Context, tenant, session, turn string) ([]FunctionCall, error) { + calls := make([]FunctionCall, 0) + err := o.storage.WithFunctionTurn(ctx, tenant, session, turn, func(ctx context.Context, tx FunctionTx, current Turn) error { + pending, err := tx.LoadPendingFunctionCalls(ctx, current.ID) + calls = append(calls, pending...) + return err + }) + if err != nil { + return nil, err + } + return calls, nil +} + +// ConfirmFunctionResult records that the executor applied the call's saved +// result, not that the tool succeeded, and journals the Turn's remaining +// required actions. A repeated receipt changes nothing. +func (o *ExecutionOperations) ConfirmFunctionResult(ctx context.Context, tenant, session, turn, call string) error { + if !validFunctionIdentity(call) { + return ErrInvalidInput + } + return o.storage.WithFunctionTurn(ctx, tenant, session, turn, func(ctx context.Context, tx FunctionTx, current Turn) error { + stored, found, err := tx.LoadFunctionCall(ctx, current.ID, call) + if err != nil { + return err + } + if !found { + return ErrNotFound + } + apply, err := decideFunctionConfirmation(current, stored) + if err != nil || !apply { + return err + } + if err := tx.ApplyFunctionResult(ctx, current.ID, call); err != nil { + return err + } + return settleFunctionState(ctx, tx, current) + }) +} + +// settleFunctionState moves a Turn that accepts results to waiting while it +// has required actions and back to in progress once it has none, journals +// that transition, then journals the Session activity with the required +// actions. +func settleFunctionState(ctx context.Context, tx FunctionTx, turn Turn) error { + actions, err := LoadRequiredActions(ctx, tx, turn) + if err != nil { + return err + } + status := TurnInProgress + if len(actions) > 0 { + status = TurnWaiting + } + if turn.Status != status { + if turn, err = tx.ApplyFunctionTurnStatus(ctx, turn.ID, turn.Status, status); err != nil { + return err + } + if err := tx.AppendChanges(ctx, TurnChanges(turn, false)...); err != nil { + return err + } + } + usage, err := tx.LoadUsage(ctx) + if err != nil { + return err + } + return tx.AppendChanges(ctx, ActivityChange(turn, usage, actions)) +} diff --git a/services/core/internal/sessions/execution_functions_test.go b/services/core/internal/sessions/execution_functions_test.go new file mode 100644 index 000000000..0b3ca82cd --- /dev/null +++ b/services/core/internal/sessions/execution_functions_test.go @@ -0,0 +1,247 @@ +package sessions + +import ( + "context" + "encoding/json" + "errors" + "reflect" + "testing" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" +) + +// fakeExecutionStorage is strict lease-bound Session storage: a method whose +// func is unset fails the test. +type fakeExecutionStorage struct { + t *testing.T + withFunctionTurn func(ctx context.Context, tenant, session, turn string, apply func(context.Context, FunctionTx, Turn) error) error +} + +var _ ExecutionStorage = (*fakeExecutionStorage)(nil) + +func (f *fakeExecutionStorage) WithFunctionTurn(ctx context.Context, tenant, session, turn string, apply func(context.Context, FunctionTx, Turn) error) error { + f.t.Helper() + if f.withFunctionTurn == nil { + f.t.Fatal("unexpected call to WithFunctionTurn") + } + return f.withFunctionTurn(ctx, tenant, session, turn, apply) +} + +const testTenant = "4f0d0c35-8b8e-4d7c-9d1c-1f0a5b8a2e61" + +// functionOperations serves the tenant's Turn as current through tx. +func functionOperations(t *testing.T, tx *fakeFunctionTx, current Turn) *ExecutionOperations { + t.Helper() + operations, err := NewExecutionOperations(&fakeExecutionStorage{t: t, withFunctionTurn: func(ctx context.Context, tenant, session, turn string, apply func(context.Context, FunctionTx, Turn) error) error { + if tenant != testTenant || session != testSession || turn != testTurn { + t.Fatalf("Turn %s %s %s", tenant, session, turn) + } + return apply(ctx, tx, current) + }}) + if err != nil { + t.Fatal(err) + } + return operations +} + +// unusedStorage fails the test on any storage call. +func unusedStorage(t *testing.T) *ExecutionOperations { + operations, err := NewExecutionOperations(&fakeExecutionStorage{t: t}) + if err != nil { + t.Fatal(err) + } + return operations +} + +func TestNewExecutionOperationsRequiresStorage(t *testing.T) { + if _, err := NewExecutionOperations(nil); err == nil { + t.Fatal("nil storage accepted") + } +} + +func TestDecideFunctionCall(t *testing.T) { + for name, test := range map[string]struct { + turn Turn + match FunctionCallMatch + record bool + err error + }{ + "new call": {turnWith(TurnInProgress), FunctionCallMatch{}, true, nil}, + "new call while waiting": {turnWith(TurnWaiting), FunctionCallMatch{}, true, nil}, + "retry": {turnWith(TurnWaiting), FunctionCallMatch{Recorded: true, Matches: true}, false, nil}, + "retry after the Turn end": {turnWith(TurnCompleted), FunctionCallMatch{Recorded: true, Matches: true}, false, nil}, + "changed call": {turnWith(TurnWaiting), FunctionCallMatch{Recorded: true}, false, ErrIdempotencyConflict}, + "queued Turn": {turnWith(TurnQueued), FunctionCallMatch{}, false, ErrTurnConflict}, + "cancelling Turn": {cancelling(TurnInProgress), FunctionCallMatch{}, false, ErrTurnConflict}, + "ended Turn": {turnWith(TurnCancelled), FunctionCallMatch{}, false, ErrTurnConflict}, + } { + record, err := decideFunctionCall(test.turn, test.match) + if record != test.record || !errors.Is(err, test.err) || (test.err == nil && err != nil) { + t.Errorf("%s: %v %v", name, record, err) + } + } +} + +func TestDecideFunctionConfirmation(t *testing.T) { + submitted := FunctionCall{CallID: "call", Result: json.RawMessage(`{"success":true}`)} + applied := submitted + applied.Applied = true + for name, test := range map[string]struct { + turn Turn + call FunctionCall + apply bool + err error + }{ + "submitted": {turnWith(TurnWaiting), submitted, true, nil}, + "already applied": {turnWith(TurnWaiting), applied, false, nil}, + "applied after the end": {turnWith(TurnCompleted), applied, false, nil}, + "no result": {turnWith(TurnWaiting), FunctionCall{CallID: "call"}, false, ErrTurnConflict}, + "cancelling": {cancelling(TurnWaiting), submitted, false, ErrTurnConflict}, + "ended": {turnWith(TurnFailed), submitted, false, ErrTurnConflict}, + } { + apply, err := decideFunctionConfirmation(test.turn, test.call) + if apply != test.apply || !errors.Is(err, test.err) || (test.err == nil && err != nil) { + t.Errorf("%s: %v %v", name, apply, err) + } + } +} + +func TestRecordFunctionCall(t *testing.T) { + usage := json.RawMessage(`{"input_tokens":1}`) + t.Run("first call moves the Turn to waiting", func(t *testing.T) { + running, waiting := turnWith(TurnInProgress), turnWith(TurnWaiting) + var changes []SessionChange + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t, loadUsage: returns(usage), appendChanges: collect(&changes)}, + matchFunctionCall: returns(FunctionCallMatch{}), createFunctionCall: done, + loadPendingFunctionCalls: returns([]FunctionCall{testCall}), applyFunctionTurnStatus: returns(waiting)} + if err := functionOperations(t, f, running).RecordFunctionCall(t.Context(), testTenant, testSession, testTurn, testCall); err != nil { + t.Fatal(err) + } + assertCalls(t, f.fakeTx, "MatchFunctionCall "+testTurn+" call", "CreateFunctionCall "+testTurn+" call", "LoadPendingFunctionCalls "+testTurn, + "ApplyFunctionTurnStatus "+testTurn+" in_progress waiting", "AppendChanges ", "LoadUsage", "AppendChanges agent.session.requires_action") + actions, _ := LoadRequiredActions(t.Context(), &fakeFunctionTx{fakeTx: &fakeTx{t: t}, loadPendingFunctionCalls: returns([]FunctionCall{testCall})}, waiting) + if !reflect.DeepEqual(changes, []SessionChange{ActivityChange(waiting, usage, actions)}) { + t.Fatalf("changes %+v", changes) + } + }) + t.Run("another call keeps the Turn waiting", func(t *testing.T) { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t, loadUsage: returns(usage), appendChanges: func([]SessionChange) error { return nil }}, + matchFunctionCall: returns(FunctionCallMatch{}), createFunctionCall: done, loadPendingFunctionCalls: returns([]FunctionCall{testCall, testCall})} + if err := functionOperations(t, f, turnWith(TurnWaiting)).RecordFunctionCall(t.Context(), testTenant, testSession, testTurn, testCall); err != nil { + t.Fatal(err) + } + assertCalls(t, f.fakeTx, "MatchFunctionCall "+testTurn+" call", "CreateFunctionCall "+testTurn+" call", "LoadPendingFunctionCalls "+testTurn, + "LoadUsage", "AppendChanges agent.session.requires_action") + }) + t.Run("retry records nothing", func(t *testing.T) { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, matchFunctionCall: returns(FunctionCallMatch{Recorded: true, Matches: true})} + if err := functionOperations(t, f, turnWith(TurnWaiting)).RecordFunctionCall(t.Context(), testTenant, testSession, testTurn, testCall); err != nil { + t.Fatal(err) + } + assertCalls(t, f.fakeTx, "MatchFunctionCall "+testTurn+" call") + }) + t.Run("conflicts record nothing", func(t *testing.T) { + for name, test := range map[string]struct { + turn Turn + match FunctionCallMatch + want error + }{ + "changed call": {turnWith(TurnWaiting), FunctionCallMatch{Recorded: true}, ErrIdempotencyConflict}, + "queued Turn": {turnWith(TurnQueued), FunctionCallMatch{}, ErrTurnConflict}, + } { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, matchFunctionCall: returns(test.match)} + if err := functionOperations(t, f, test.turn).RecordFunctionCall(t.Context(), testTenant, testSession, testTurn, testCall); !errors.Is(err, test.want) { + t.Errorf("%s: %v", name, err) + } + assertCalls(t, f.fakeTx, "MatchFunctionCall "+testTurn+" call") + } + }) + t.Run("an invalid call reaches no storage", func(t *testing.T) { + call := testCall + call.Name = "" + if err := unusedStorage(t).RecordFunctionCall(t.Context(), testTenant, testSession, testTurn, call); !errors.Is(err, ErrInvalidInput) { + t.Fatal(err) + } + }) + t.Run("storage errors", func(t *testing.T) { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, matchFunctionCall: returns(FunctionCallMatch{}), createFunctionCall: func() error { return errStorage }} + if err := functionOperations(t, f, turnWith(TurnInProgress)).RecordFunctionCall(t.Context(), testTenant, testSession, testTurn, testCall); !errors.Is(err, errStorage) { + t.Fatal(err) + } + }) +} + +func TestConfirmFunctionResult(t *testing.T) { + usage := json.RawMessage(`{"input_tokens":1}`) + submitted := FunctionCall{CallID: "call", Result: json.RawMessage(`{"success":true}`)} + t.Run("last receipt resumes the Turn", func(t *testing.T) { + waiting, running := turnWith(TurnWaiting), turnWith(TurnInProgress) + var changes []SessionChange + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t, loadUsage: returns(usage), appendChanges: collect(&changes)}, + loadFunctionCall: func() (FunctionCall, bool, error) { return submitted, true, nil }, applyFunctionResult: done, + loadPendingFunctionCalls: returns([]FunctionCall{}), applyFunctionTurnStatus: returns(running)} + if err := functionOperations(t, f, waiting).ConfirmFunctionResult(t.Context(), testTenant, testSession, testTurn, "call"); err != nil { + t.Fatal(err) + } + assertCalls(t, f.fakeTx, "LoadFunctionCall "+testTurn+" call", "ApplyFunctionResult "+testTurn+" call", "LoadPendingFunctionCalls "+testTurn, + "ApplyFunctionTurnStatus "+testTurn+" waiting in_progress", "AppendChanges agent.session.turn.in_progress", "LoadUsage", "AppendChanges agent.session.in_progress") + want := append(TurnChanges(running, false), ActivityChange(running, usage, []v1.FunctionCallAction{})) + if !reflect.DeepEqual(changes, want) { + t.Fatalf("changes %+v", changes) + } + }) + t.Run("repeated receipt records nothing", func(t *testing.T) { + applied := submitted + applied.Applied = true + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, loadFunctionCall: func() (FunctionCall, bool, error) { return applied, true, nil }} + if err := functionOperations(t, f, turnWith(TurnCompleted)).ConfirmFunctionResult(t.Context(), testTenant, testSession, testTurn, "call"); err != nil { + t.Fatal(err) + } + assertCalls(t, f.fakeTx, "LoadFunctionCall "+testTurn+" call") + }) + t.Run("rejections record nothing", func(t *testing.T) { + for name, test := range map[string]struct { + turn Turn + call FunctionCall + found bool + want error + }{ + "missing call": {turnWith(TurnWaiting), FunctionCall{}, false, ErrNotFound}, + "no result": {turnWith(TurnWaiting), FunctionCall{CallID: "call"}, true, ErrTurnConflict}, + "cancelling": {cancelling(TurnWaiting), submitted, true, ErrTurnConflict}, + } { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, loadFunctionCall: func() (FunctionCall, bool, error) { return test.call, test.found, nil }} + if err := functionOperations(t, f, test.turn).ConfirmFunctionResult(t.Context(), testTenant, testSession, testTurn, "call"); !errors.Is(err, test.want) { + t.Errorf("%s: %v", name, err) + } + assertCalls(t, f.fakeTx, "LoadFunctionCall "+testTurn+" call") + } + }) + t.Run("an invalid call reaches no storage", func(t *testing.T) { + if err := unusedStorage(t).ConfirmFunctionResult(t.Context(), testTenant, testSession, testTurn, " "); !errors.Is(err, ErrInvalidInput) { + t.Fatal(err) + } + }) +} + +func TestPendingFunctionCalls(t *testing.T) { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, loadPendingFunctionCalls: returns([]FunctionCall{testCall})} + calls, err := functionOperations(t, f, turnWith(TurnWaiting)).PendingFunctionCalls(t.Context(), testTenant, testSession, testTurn) + if err != nil || !reflect.DeepEqual(calls, []FunctionCall{testCall}) { + t.Fatal(calls, err) + } + assertCalls(t, f.fakeTx, "LoadPendingFunctionCalls "+testTurn) + + // Pending calls are read whatever the Turn's status; the read itself + // selects none for a cancelling or ended Turn. + none := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, loadPendingFunctionCalls: returns([]FunctionCall(nil))} + calls, err = functionOperations(t, none, turnWith(TurnCompleted)).PendingFunctionCalls(t.Context(), testTenant, testSession, testTurn) + if err != nil || calls == nil || len(calls) != 0 { + t.Fatal(calls, err) + } + + failing := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, loadPendingFunctionCalls: func() ([]FunctionCall, error) { return []FunctionCall{testCall}, errStorage }} + if calls, err := functionOperations(t, failing, turnWith(TurnWaiting)).PendingFunctionCalls(t.Context(), testTenant, testSession, testTurn); !errors.Is(err, errStorage) || calls != nil { + t.Fatal(calls, err) + } +} diff --git a/services/core/internal/sessions/functions.go b/services/core/internal/sessions/functions.go new file mode 100644 index 000000000..df4832f21 --- /dev/null +++ b/services/core/internal/sessions/functions.go @@ -0,0 +1,159 @@ +package sessions + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" +) + +// maxFunctionArguments bounds the arguments of a recorded function call. +const maxFunctionArguments = 512 * 1024 + +// validFunctionIdentity checks a function call's public or executor identity or +// its name. +func validFunctionIdentity(id string) bool { return strings.TrimSpace(id) != "" && len(id) <= 512 } + +// validateFunctionCall checks a call an executor reported before it is +// recorded: identities and name within bounds, JSON arguments within bounds and +// no result or application receipt yet. +func validateFunctionCall(call FunctionCall) error { + if !validFunctionIdentity(call.CallID) || !validFunctionIdentity(call.ExecutorCallID) || !validFunctionIdentity(call.Name) || + len(call.Arguments) > maxFunctionArguments || !json.Valid(call.Arguments) || call.Result != nil || call.Applied { + return ErrInvalidInput + } + return nil +} + +// ParseFunctionResultInput decodes a tool_result input payload: a Turn ID, a +// call identity and a result object, which it normalizes. It does not check +// that the Turn or call exists; AdmitFunctionResult resolves them. +func ParseFunctionResultInput(raw json.RawMessage) (FunctionResultInput, error) { + var input FunctionResultInput + if json.Unmarshal(raw, &input) != nil || input.TurnID == "" || !validFunctionIdentity(input.CallID) || len(input.Result) == 0 { + return input, ErrInvalidInput + } + result, err := jsonobject.Normalize(input.Result) + if err != nil { + return input, fmt.Errorf("%w: %w", ErrInvalidInput, err) + } + input.Result = result + return input, nil +} + +// acceptsFunctionResults reports whether a Turn takes new function calls, +// results and application receipts: it is in progress or waiting, and no +// cancellation was requested. +func acceptsFunctionResults(turn Turn) bool { + return (turn.Status == TurnInProgress || turn.Status == TurnWaiting) && turn.CancelRequestedAt.IsZero() +} + +// FunctionResultMatch is what a Turn's stored call shows about a submitted +// result. +type FunctionResultMatch struct { + // Recorded reports that the Turn has the call. + Recorded bool + // Submitted reports that the call already has a result. + Submitted bool + // Matches reports that the saved result equals the submitted one. + Matches bool +} + +// decideFunctionResult decides whether to store a submitted result for a call +// the Turn has. An identical retry of a saved result stores nothing, even after +// the Turn ended; a different result conflicts with the saved one. A new result +// needs a Turn that accepts results. +func decideFunctionResult(turn Turn, match FunctionResultMatch) (bool, error) { + if match.Submitted { + if !match.Matches { + return false, ErrFunctionResultConflict + } + return false, nil + } + if !acceptsFunctionResults(turn) { + return false, ErrTurnConflict + } + return true, nil +} + +// FunctionResultTx is the Session transaction a function result is admitted in. +type FunctionResultTx interface { + // FindResultTurn reads the Session's Turn that a result names and reports + // whether the Session has it. A malformed ID names no Turn. + FindResultTurn(ctx context.Context, turn string) (Turn, bool, error) + // MatchFunctionResult compares result with the Turn's call. + MatchFunctionResult(ctx context.Context, turn, call string, result json.RawMessage) (FunctionResultMatch, error) + // SubmitFunctionResult saves the result of the Turn's call. + SubmitFunctionResult(ctx context.Context, turn, call string, result json.RawMessage) error + // HasFunctionCall reports whether any Turn of the Session has the call. + HasFunctionCall(ctx context.Context, call string) (bool, error) +} + +// AdmitFunctionResult saves a submitted function result for its Turn's call +// and returns that Turn, which the result's input joins. It never creates a +// Turn. A result whose Turn does not have the call is +// ErrFunctionCallTurnMismatch when another Turn of the Session has it and +// ErrUnknownFunctionCall otherwise; the answer depends only on the Session's +// own calls. +func AdmitFunctionResult(ctx context.Context, tx FunctionResultTx, input FunctionResultInput) (Turn, error) { + turn, found, err := tx.FindResultTurn(ctx, input.TurnID) + if err != nil { + return Turn{}, err + } + if found { + match, err := tx.MatchFunctionResult(ctx, turn.ID, input.CallID, input.Result) + if err != nil { + return Turn{}, err + } + if match.Recorded { + submit, err := decideFunctionResult(turn, match) + if err != nil { + return Turn{}, err + } + if submit { + if err := tx.SubmitFunctionResult(ctx, turn.ID, input.CallID, input.Result); err != nil { + return Turn{}, err + } + } + return turn, nil + } + } + elsewhere, err := tx.HasFunctionCall(ctx, input.CallID) + if err != nil { + return Turn{}, err + } + if elsewhere { + return Turn{}, ErrFunctionCallTurnMismatch + } + return Turn{}, ErrUnknownFunctionCall +} + +// PendingFunctionCallsTx reads a Turn's function calls that wait for a result +// or its application receipt. +type PendingFunctionCallsTx interface { + // LoadPendingFunctionCalls reads the Turn's calls without an application + // receipt, in creation order, while the Turn is in progress or waiting and + // no cancellation was requested, and none otherwise. + LoadPendingFunctionCalls(ctx context.Context, turn string) ([]FunctionCall, error) +} + +// LoadRequiredActions returns the function calls a Turn requires the caller to +// answer as public required actions: its pending calls while it accepts +// results, and none otherwise. +func LoadRequiredActions(ctx context.Context, tx PendingFunctionCallsTx, turn Turn) ([]v1.FunctionCallAction, error) { + actions := make([]v1.FunctionCallAction, 0) + if !acceptsFunctionResults(turn) { + return actions, nil + } + calls, err := tx.LoadPendingFunctionCalls(ctx, turn.ID) + if err != nil { + return nil, err + } + for _, call := range calls { + actions = append(actions, v1.FunctionCallAction{Type: "function_call", CallID: call.CallID, Name: call.Name, TurnID: turn.ID, Arguments: call.Arguments}) + } + return actions, nil +} diff --git a/services/core/internal/sessions/functions_test.go b/services/core/internal/sessions/functions_test.go new file mode 100644 index 000000000..1b1a33f5b --- /dev/null +++ b/services/core/internal/sessions/functions_test.go @@ -0,0 +1,298 @@ +package sessions + +import ( + "context" + "encoding/json" + "errors" + "reflect" + "strings" + "testing" + "time" +) + +// fakeFunctionTx is fakeTx with the function-call methods, equally strict: it +// shares fakeTx's call log, and a method whose func is unset fails the test. +type fakeFunctionTx struct { + *fakeTx + + loadPendingFunctionCalls func() ([]FunctionCall, error) + findResultTurn func() (Turn, bool, error) + matchFunctionResult func() (FunctionResultMatch, error) + submitFunctionResult func() error + hasFunctionCall func() (bool, error) + matchFunctionCall func() (FunctionCallMatch, error) + createFunctionCall func() error + loadFunctionCall func() (FunctionCall, bool, error) + applyFunctionResult func() error + applyFunctionTurnStatus func() (Turn, error) +} + +var ( + _ FunctionResultTx = (*fakeFunctionTx)(nil) + _ FunctionTx = (*fakeFunctionTx)(nil) +) + +func (f *fakeFunctionTx) LoadPendingFunctionCalls(_ context.Context, turn string) ([]FunctionCall, error) { + f.record("LoadPendingFunctionCalls", f.loadPendingFunctionCalls != nil, turn) + return f.loadPendingFunctionCalls() +} + +func (f *fakeFunctionTx) FindResultTurn(_ context.Context, turn string) (Turn, bool, error) { + f.record("FindResultTurn", f.findResultTurn != nil, turn) + return f.findResultTurn() +} + +func (f *fakeFunctionTx) MatchFunctionResult(_ context.Context, turn, call string, result json.RawMessage) (FunctionResultMatch, error) { + f.record("MatchFunctionResult", f.matchFunctionResult != nil, turn, call, string(result)) + return f.matchFunctionResult() +} + +func (f *fakeFunctionTx) SubmitFunctionResult(_ context.Context, turn, call string, result json.RawMessage) error { + f.record("SubmitFunctionResult", f.submitFunctionResult != nil, turn, call, string(result)) + return f.submitFunctionResult() +} + +func (f *fakeFunctionTx) HasFunctionCall(_ context.Context, call string) (bool, error) { + f.record("HasFunctionCall", f.hasFunctionCall != nil, call) + return f.hasFunctionCall() +} + +func (f *fakeFunctionTx) MatchFunctionCall(_ context.Context, turn string, call FunctionCall) (FunctionCallMatch, error) { + f.record("MatchFunctionCall", f.matchFunctionCall != nil, turn, call.CallID) + return f.matchFunctionCall() +} + +func (f *fakeFunctionTx) CreateFunctionCall(_ context.Context, turn string, call FunctionCall) error { + f.record("CreateFunctionCall", f.createFunctionCall != nil, turn, call.CallID) + return f.createFunctionCall() +} + +func (f *fakeFunctionTx) LoadFunctionCall(_ context.Context, turn, call string) (FunctionCall, bool, error) { + f.record("LoadFunctionCall", f.loadFunctionCall != nil, turn, call) + return f.loadFunctionCall() +} + +func (f *fakeFunctionTx) ApplyFunctionResult(_ context.Context, turn, call string) error { + f.record("ApplyFunctionResult", f.applyFunctionResult != nil, turn, call) + return f.applyFunctionResult() +} + +func (f *fakeFunctionTx) ApplyFunctionTurnStatus(_ context.Context, turn, expected, status string) (Turn, error) { + f.record("ApplyFunctionTurnStatus", f.applyFunctionTurnStatus != nil, turn, expected, status) + return f.applyFunctionTurnStatus() +} + +func cancelling(status string) Turn { + turn := turnWith(status) + turn.CancelRequestedAt = time.Unix(1700000001, 0) + return turn +} + +var testCall = FunctionCall{CallID: "call", ExecutorCallID: "native-call", Name: "lookup", Arguments: json.RawMessage(`{"ticket":9007199254740993}`)} + +func TestValidateFunctionCall(t *testing.T) { + for name, change := range map[string]func(*FunctionCall){ + "blank call": func(c *FunctionCall) { c.CallID = " " }, + "long call": func(c *FunctionCall) { c.CallID = strings.Repeat("c", 513) }, + "blank executor": func(c *FunctionCall) { c.ExecutorCallID = "" }, + "blank name": func(c *FunctionCall) { c.Name = "\t" }, + "long name": func(c *FunctionCall) { c.Name = strings.Repeat("n", 513) }, + "invalid arguments": func(c *FunctionCall) { c.Arguments = json.RawMessage(`{`) }, + "missing arguments": func(c *FunctionCall) { c.Arguments = nil }, + "oversized": func(c *FunctionCall) { + c.Arguments = json.RawMessage(`"` + strings.Repeat("a", maxFunctionArguments-1) + `"`) + }, + "result": func(c *FunctionCall) { c.Result = json.RawMessage(`{}`) }, + "applied": func(c *FunctionCall) { c.Applied = true }, + "empty result slice": func(c *FunctionCall) { c.Result = json.RawMessage{} }, + } { + call := testCall + change(&call) + if err := validateFunctionCall(call); !errors.Is(err, ErrInvalidInput) { + t.Errorf("%s: %v", name, err) + } + } + bounded := testCall + bounded.CallID, bounded.Arguments = strings.Repeat("c", 512), json.RawMessage(`"`+strings.Repeat("a", maxFunctionArguments-2)+`"`) + for _, call := range []FunctionCall{testCall, bounded} { + if err := validateFunctionCall(call); err != nil { + t.Error(err) + } + } +} + +func TestParseFunctionResultInput(t *testing.T) { + input, err := ParseFunctionResultInput(json.RawMessage(`{"turn_id":"not-a-uuid","call_id":"call","result":{ "success" : true, "output":[{"type":"input_text","text":"a"}] }}`)) + if err != nil || input.TurnID != "not-a-uuid" || input.CallID != "call" || string(input.Result) != `{"output":[{"text":"a","type":"input_text"}],"success":true}` { + t.Fatalf("%+v %v", input, err) + } + for name, raw := range map[string]string{ + "malformed": `{`, + "missing turn": `{"call_id":"call","result":{}}`, + "blank call": `{"turn_id":"turn","call_id":" ","result":{}}`, + "long call": `{"turn_id":"turn","call_id":"` + strings.Repeat("c", 513) + `","result":{}}`, + "missing result": `{"turn_id":"turn","call_id":"call"}`, + "null result": `{"turn_id":"turn","call_id":"call","result":null}`, + "array result": `{"turn_id":"turn","call_id":"call","result":[]}`, + "string result": `{"turn_id":"turn","call_id":"call","result":"ok"}`, + } { + if _, err := ParseFunctionResultInput(json.RawMessage(raw)); !errors.Is(err, ErrInvalidInput) { + t.Errorf("%s: %v", name, err) + } + } +} + +func TestAcceptsFunctionResults(t *testing.T) { + for name, test := range map[string]struct { + turn Turn + want bool + }{ + "queued": {turnWith(TurnQueued), false}, + "in progress": {turnWith(TurnInProgress), true}, + "waiting": {turnWith(TurnWaiting), true}, + "cancelling": {cancelling(TurnInProgress), false}, + "cancelling and waiting": {cancelling(TurnWaiting), false}, + "completed": {turnWith(TurnCompleted), false}, + "failed": {turnWith(TurnFailed), false}, + "cancelled": {turnWith(TurnCancelled), false}, + } { + if got := acceptsFunctionResults(test.turn); got != test.want { + t.Errorf("%s: %v", name, got) + } + } +} + +func TestDecideFunctionResult(t *testing.T) { + saved, same := FunctionResultMatch{Recorded: true, Submitted: true}, FunctionResultMatch{Recorded: true, Submitted: true, Matches: true} + for name, test := range map[string]struct { + turn Turn + match FunctionResultMatch + store bool + err error + }{ + "new result": {turnWith(TurnWaiting), FunctionResultMatch{Recorded: true}, true, nil}, + "identical retry": {turnWith(TurnWaiting), same, false, nil}, + "identical retry after end": {turnWith(TurnCompleted), same, false, nil}, + "different result": {turnWith(TurnWaiting), saved, false, ErrFunctionResultConflict}, + "different result after end": {turnWith(TurnCancelled), saved, false, ErrFunctionResultConflict}, + "new result while cancelling": {cancelling(TurnWaiting), FunctionResultMatch{Recorded: true}, false, ErrTurnConflict}, + "new result after the Turn end": {turnWith(TurnFailed), FunctionResultMatch{Recorded: true}, false, ErrTurnConflict}, + } { + store, err := decideFunctionResult(test.turn, test.match) + if store != test.store || !errors.Is(err, test.err) || (test.err == nil && err != nil) { + t.Errorf("%s: %v %v", name, store, err) + } + } +} + +func TestAdmitFunctionResult(t *testing.T) { + input := FunctionResultInput{TurnID: testTurn, CallID: "call", Result: json.RawMessage(`{"success":true}`)} + waiting := turnWith(TurnWaiting) + find, match, submit := "FindResultTurn "+testTurn, `MatchFunctionResult `+testTurn+` call {"success":true}`, `SubmitFunctionResult `+testTurn+` call {"success":true}` + t.Run("saves a new result", func(t *testing.T) { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, findResultTurn: found(waiting), matchFunctionResult: returns(FunctionResultMatch{Recorded: true}), submitFunctionResult: done} + turn, err := AdmitFunctionResult(t.Context(), f, input) + if err != nil || !reflect.DeepEqual(turn, waiting) { + t.Fatal(turn, err) + } + assertCalls(t, f.fakeTx, find, match, submit) + }) + t.Run("identical retry keeps the saved result", func(t *testing.T) { + completed := turnWith(TurnCompleted) + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, findResultTurn: found(completed), matchFunctionResult: returns(FunctionResultMatch{Recorded: true, Submitted: true, Matches: true})} + if turn, err := AdmitFunctionResult(t.Context(), f, input); err != nil || turn.ID != testTurn { + t.Fatal(turn, err) + } + assertCalls(t, f.fakeTx, find, match) + }) + t.Run("decision errors", func(t *testing.T) { + for name, test := range map[string]struct { + turn Turn + match FunctionResultMatch + want error + }{ + "changed result": {waiting, FunctionResultMatch{Recorded: true, Submitted: true}, ErrFunctionResultConflict}, + "Turn is cancelling": {cancelling(TurnWaiting), FunctionResultMatch{Recorded: true}, ErrTurnConflict}, + "Turn already finished": {turnWith(TurnCompleted), FunctionResultMatch{Recorded: true}, ErrTurnConflict}, + } { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, findResultTurn: found(test.turn), matchFunctionResult: returns(test.match)} + if _, err := AdmitFunctionResult(t.Context(), f, input); !errors.Is(err, test.want) { + t.Errorf("%s: %v", name, err) + } + assertCalls(t, f.fakeTx, find, match) + } + }) + t.Run("classifies a call the Turn does not have", func(t *testing.T) { + for name, test := range map[string]struct { + turnFound, elsewhere bool + want error + }{ + "call of another Turn": {true, true, ErrFunctionCallTurnMismatch}, + "unknown call": {true, false, ErrUnknownFunctionCall}, + "unknown Turn, call elsewhere": {false, true, ErrFunctionCallTurnMismatch}, + "unknown Turn and call": {false, false, ErrUnknownFunctionCall}, + } { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, hasFunctionCall: returns(test.elsewhere), matchFunctionResult: returns(FunctionResultMatch{}), + findResultTurn: func() (Turn, bool, error) { return waiting, test.turnFound, nil }} + if _, err := AdmitFunctionResult(t.Context(), f, input); !errors.Is(err, test.want) { + t.Errorf("%s: %v", name, err) + } + want := []string{find, match, "HasFunctionCall call"} + if !test.turnFound { + want = []string{find, "HasFunctionCall call"} + } + assertCalls(t, f.fakeTx, want...) + } + }) + t.Run("storage errors", func(t *testing.T) { + for name, f := range map[string]*fakeFunctionTx{ + "find": {fakeTx: &fakeTx{t: t}, findResultTurn: func() (Turn, bool, error) { return Turn{}, false, errStorage }}, + "match": {fakeTx: &fakeTx{t: t}, findResultTurn: found(waiting), matchFunctionResult: func() (FunctionResultMatch, error) { return FunctionResultMatch{}, errStorage }}, + "submit": {fakeTx: &fakeTx{t: t}, findResultTurn: found(waiting), matchFunctionResult: returns(FunctionResultMatch{Recorded: true}), submitFunctionResult: func() error { return errStorage }}, + "classify": {fakeTx: &fakeTx{t: t}, findResultTurn: func() (Turn, bool, error) { return Turn{}, false, nil }, + hasFunctionCall: func() (bool, error) { return false, errStorage }}, + } { + if _, err := AdmitFunctionResult(t.Context(), f, input); !errors.Is(err, errStorage) { + t.Errorf("%s: %v", name, err) + } + } + }) +} + +func TestLoadRequiredActions(t *testing.T) { + calls := []FunctionCall{testCall, {CallID: "second", ExecutorCallID: "native-second", Name: "search", Arguments: json.RawMessage(`{}`), Result: json.RawMessage(`{"success":true}`)}} + for _, status := range []string{TurnInProgress, TurnWaiting} { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, loadPendingFunctionCalls: returns(calls)} + actions, err := LoadRequiredActions(t.Context(), f, turnWith(status)) + if err != nil || len(actions) != 2 { + t.Fatal(actions, err) + } + assertCalls(t, f.fakeTx, "LoadPendingFunctionCalls "+testTurn) + for i, action := range actions { + if action.Type != "function_call" || action.CallID != calls[i].CallID || action.Name != calls[i].Name || action.TurnID != testTurn { + t.Fatalf("action %d: %+v", i, action) + } + // Arguments keep their stored bytes, so large integers keep their precision. + if raw, ok := action.Arguments.(json.RawMessage); !ok || string(raw) != string(calls[i].Arguments) { + t.Fatalf("arguments %d: %#v", i, action.Arguments) + } + } + } + for _, turn := range []Turn{turnWith(TurnQueued), cancelling(TurnWaiting), turnWith(TurnCompleted)} { + f := &fakeFunctionTx{fakeTx: &fakeTx{t: t}} + actions, err := LoadRequiredActions(t.Context(), f, turn) + if err != nil || actions == nil || len(actions) != 0 { + t.Fatal(turn.Status, actions, err) + } + assertCalls(t, f.fakeTx) + } + failing := &fakeFunctionTx{fakeTx: &fakeTx{t: t}, loadPendingFunctionCalls: func() ([]FunctionCall, error) { return nil, errStorage }} + if _, err := LoadRequiredActions(t.Context(), failing, turnWith(TurnWaiting)); !errors.Is(err, errStorage) { + t.Fatal(err) + } +} + +// found is a fake lookup that finds turn. +func found(turn Turn) func() (Turn, bool, error) { + return func() (Turn, bool, error) { return turn, true, nil } +} diff --git a/services/core/internal/store/archive_cancellation_test.go b/services/core/internal/store/archive_cancellation_test.go index 93aeb7fb5..b891ebd2f 100644 --- a/services/core/internal/store/archive_cancellation_test.go +++ b/services/core/internal/store/archive_cancellation_test.go @@ -104,7 +104,8 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { if err != nil { t.Fatal(err) } - h := &dispatchHarness{t: t, s: s, db: db, lease: leased.Lease, tenant: project.TenantID, session: session, conn: conn, registry: registry, d: &execution.Dispatcher{Store: writer, Registry: registry, Observer: modelconfigurationpg.New(pgunit.NewPool(db.pool), db.cipher), Sessions: sessionService, SessionsReader: sessionStore}} + h := &dispatchHarness{t: t, s: s, db: db, lease: leased.Lease, owned: &leased, tenant: project.TenantID, session: session, conn: conn, registry: registry, d: &execution.Dispatcher{Registry: registry, Observer: modelconfigurationpg.New(pgunit.NewPool(db.pool), db.cipher), Sessions: sessionService, SessionsReader: sessionStore}} + h.d = h.bound() capabilities := workerEnvironmentCapabilities() capabilities.FunctionTools = proto.CapabilitySupported h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: capabilities}}}) diff --git a/services/core/internal/store/claude_execution_test.go b/services/core/internal/store/claude_execution_test.go index 95cfaddfa..2d5503acb 100644 --- a/services/core/internal/store/claude_execution_test.go +++ b/services/core/internal/store/claude_execution_test.go @@ -123,14 +123,14 @@ func TestClaudeDispatcherRejectsUnsupportedConfigurationBeforeClaim(t *testing.T func TestClaudeInvalidImageResultRejectsWholeBatchBeforePersistence(t *testing.T) { h := newDispatchHarness(t) claudeSession(t, h, functionConfiguration, false) - worker := startWorker(t, t.Context(), h.db, h.d) + worker := startOwnedWorker(t, t.Context(), h.db, h.d, h.owner()) defer func() { ctx, cancel := context.WithCancel(context.Background()); cancel(); _ = worker.Run(ctx) }() input := h.message("start", "Run") if _, err := h.s.TransitionTurn(t.Context(), h.tenant, h.session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } call := sessions.FunctionCall{CallID: "public-call", ExecutorCallID: "native-call", Name: "lookup_ticket", Arguments: json.RawMessage(`{"ticket":"42"}`)} - if err := h.s.RecordFunctionCall(t.Context(), h.tenant, h.session.ID, input.TurnID, call); err != nil { + if err := h.owner().Sessions.RecordFunctionCall(t.Context(), h.tenant, h.session.ID, input.TurnID, call); err != nil { t.Fatal(err) } result := func(raw string) sessions.Input { @@ -144,7 +144,7 @@ func TestClaudeInvalidImageResultRejectsWholeBatchBeforePersistence(t *testing.T if _, err := worker.SubmitInputs(t.Context(), h.tenant, h.session.ID, "batch", batch); !errors.Is(err, sessions.ErrInvalidInput) { t.Fatal(err) } - saved, err := h.s.GetFunctionCall(t.Context(), h.tenant, h.session.ID, input.TurnID, call.CallID) + saved, err := store.FixtureFunctionCall(t.Context(), h.db.pool, h.tenant, h.session.ID, input.TurnID, call.CallID) if err != nil || saved.Result != nil || saved.Applied { t.Fatal(saved, err) } diff --git a/services/core/internal/store/claude_mcp_test.go b/services/core/internal/store/claude_mcp_test.go index 92d4b6ea7..ccab5314a 100644 --- a/services/core/internal/store/claude_mcp_test.go +++ b/services/core/internal/store/claude_mcp_test.go @@ -43,7 +43,7 @@ func TestClaudeMCPWaitsForCapableRuntime(t *testing.T) { } ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker := startWorker(t, ctx, h.db, h.d) + worker := startOwnedWorker(t, ctx, h.db, h.d, h.owner()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { diff --git a/services/core/internal/store/dispatch_test.go b/services/core/internal/store/dispatch_test.go index 274948591..0e36034d5 100644 --- a/services/core/internal/store/dispatch_test.go +++ b/services/core/internal/store/dispatch_test.go @@ -15,6 +15,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/modelconfigurationpg" + "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/runtime" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" @@ -33,6 +34,7 @@ type dispatchHarness struct { db fixtureDB d *execution.Dispatcher lease execution.Ownership // held by tests that run execution operations without a Worker + owned *execution.Owner // the Owner that bound binds, acquired on first use tenant string session sessions.Session device sessions.ExecutionDevice @@ -175,9 +177,34 @@ type runResult struct { err error } +// owner returns the harness's execution Owner. The first call acquires the +// execution lease, which closes when the test ends. A test that also starts a +// Worker hands it this Owner through startOwnedWorker. +func (h *dispatchHarness) owner() execution.Owner { + h.t.Helper() + if h.owned == nil { + owner := executionOwner(h.t, h.db, h.s) + h.owned = &owner + } + return *h.owned +} + +// bound returns h.d bound to the harness's execution Owner, as StartWorker +// binds a Worker's Dispatcher. +func (h *dispatchHarness) bound() *execution.Dispatcher { + h.t.Helper() + d, err := h.d.Bind(h.owner()) + if err != nil { + h.t.Fatal(err) + } + return d +} + func (h *dispatchHarness) run(ctx context.Context, turn string) <-chan runResult { + h.t.Helper() + d := h.bound() out := make(chan runResult, 1) - go func() { result, err := h.d.Run(ctx, h.tenant, h.session.ID, turn); out <- runResult{result, err} }() + go func() { result, err := d.Run(ctx, h.tenant, h.session.ID, turn); out <- runResult{result, err} }() return out } @@ -206,10 +233,10 @@ func TestExecutionDispatchSteeringAndNativeContinuity(t *testing.T) { if inputTextForTest(t, prompt.Input) != "Initial input" || prompt.ConversationID != h.session.ID || prompt.AgentOptions["model"] != "test-model" || prompt.AgentOptions["system_prompt"] != "Keep this instruction." { t.Fatalf("wrong resolved request: %+v", prompt) } - if _, err := h.d.Run(ctx, uuid.NewString(), h.session.ID, first.TurnID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := h.bound().Run(ctx, uuid.NewString(), h.session.ID, first.TurnID); !errors.Is(err, sessions.ErrNotFound) { t.Fatalf("foreign execution: %v", err) } - if _, err := h.d.Run(ctx, h.tenant, h.session.ID, first.TurnID); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := h.bound().Run(ctx, h.tenant, h.session.ID, first.TurnID); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatalf("duplicate execution: %v", err) } second := h.message("second", "Follow-up input") @@ -232,9 +259,14 @@ func TestExecutionDispatchSteeringAndNativeContinuity(t *testing.T) { if outcome.AppliedThrough != second.Sequence || outcome.Done.Usage.InputTokens != 7 { t.Fatalf("missing result: %+v", outcome) } + // Restart Core: a new Store and execution owner continue the native Session. + awaitRelease := pgtest.ObserveExecutionLeaseRelease(t, h.db.pool) + if err := h.owner().Lease.Close(ctx); err != nil { + t.Fatal(err) + } + awaitRelease() newStore, db := newTestStoreDB(t) - defer db.pool.Close() - h.s, h.db = newStore, db + h.s, h.db, h.owned = newStore, db, nil h.d.Store = newStore bound, err := newStore.GetSessionExecutionBinding(ctx, h.tenant, h.session.ID) if err != nil || bound.NativeSessionID != "native-thread-1" { @@ -335,7 +367,7 @@ func TestExecutionFailureDoesNotBecomeSuccessOrReplay(t *testing.T) { } } if kind != "disconnect" { - if _, err := h.d.Run(context.Background(), h.tenant, h.session.ID, first.TurnID); !errors.Is(err, sessions.ErrTurnConflict) { + if _, err := h.bound().Run(context.Background(), h.tenant, h.session.ID, first.TurnID); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatalf("terminal replay: %v", err) } } @@ -413,7 +445,7 @@ func TestExecutionRejectsRuntimeMissingCapabilityBeforeClaim(t *testing.T) { time.Sleep(10 * time.Millisecond) } first := h.message("missing-capability", "Run") - if _, err := h.d.Run(context.Background(), h.tenant, h.session.ID, first.TurnID); err == nil || err.Error() != tc.message { + if _, err := h.bound().Run(context.Background(), h.tenant, h.session.ID, first.TurnID); err == nil || err.Error() != tc.message { t.Fatalf("Run error = %v, want %q", err, tc.message) } turn, err := h.s.GetTurn(context.Background(), h.tenant, h.session.ID, first.TurnID) diff --git a/services/core/internal/store/environment_runtime_fixture_test.go b/services/core/internal/store/environment_runtime_fixture_test.go index 501904ad7..c400fd523 100644 --- a/services/core/internal/store/environment_runtime_fixture_test.go +++ b/services/core/internal/store/environment_runtime_fixture_test.go @@ -39,7 +39,7 @@ func enrollFixtureSession(t *testing.T, s *store.Store, tenant string, session s func connectFixtureRuntime(t *testing.T, h *dispatchHarness, session sessions.Session) *dispatchHarness { t.Helper() // The Runtime shares the harness's Core, not its connection or write lock. - other := &dispatchHarness{t: h.t, s: h.s, db: h.db, lease: h.lease, d: h.d, tenant: h.tenant, session: session, registry: h.registry, url: h.url, + other := &dispatchHarness{t: h.t, s: h.s, db: h.db, lease: h.lease, owned: h.owned, d: h.d, tenant: h.tenant, session: session, registry: h.registry, url: h.url, admissions: h.admissions, environments: h.environments} other.device, other.credential = enrollFixtureSession(t, h.s, h.tenant, session) u, err := url.Parse(h.url) diff --git a/services/core/internal/store/execution_test.go b/services/core/internal/store/execution_test.go index eaecd1dbc..b036e3bf2 100644 --- a/services/core/internal/store/execution_test.go +++ b/services/core/internal/store/execution_test.go @@ -14,6 +14,7 @@ import ( "github.com/jackc/pgx/v5/pgxpool" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -31,6 +32,29 @@ func executionWriter(t *testing.T, s *Store) *Store { return NewExecution(s, lease) } +// sessionExecution builds the Session execution operations on lease, as +// cmd/server does. +func sessionExecution(t *testing.T, lease *pgunit.Lease) *sessions.ExecutionOperations { + t.Helper() + operations, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) + if err != nil { + t.Fatal(err) + } + return operations +} + +// functionExecution builds the Session execution operations on the execution +// lease of a pool of its own, so the test may close its Store's pool. +func functionExecution(t *testing.T) *sessions.ExecutionOperations { + t.Helper() + other, _ := testStore(t) + return sessionExecution(t, executionWriter(t, other).lease) +} + +func functionCallFixture(id string) sessions.FunctionCall { + return sessions.FunctionCall{CallID: id, ExecutorCallID: "native-" + id, Name: "lookup", Arguments: json.RawMessage(`{"ticket":9007199254740993}`)} +} + // executionOwnerPID finds the backend holding this database's execution lease, // matching the single-bigint key in queries/scheduling.sql. func executionOwnerPID(t *testing.T, pool *pgxpool.Pool) int32 { @@ -48,6 +72,7 @@ func executionOwnerPID(t *testing.T, pool *pgxpool.Pool) int32 { func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { s, pool := testStore(t) writer := executionWriter(t, s) + functions := sessionExecution(t, writer.lease) tenant, active := newTurnSession(t, s) input := submitMessage(t, s, tenant, active.ID, "active") transition(t, writer, tenant, active.ID, input.TurnID, sessions.TurnQueued, sessions.TurnInProgress) @@ -70,10 +95,10 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { waitInput := submitMessage(t, s, tenant, waiting.ID, "waiting") transition(t, writer, tenant, waiting.ID, waitInput.TurnID, sessions.TurnQueued, sessions.TurnInProgress) call := functionCallFixture("saved") - if err = writer.RecordFunctionCall(t.Context(), tenant, waiting.ID, waitInput.TurnID, call); err != nil { + if err = functions.RecordFunctionCall(t.Context(), tenant, waiting.ID, waitInput.TurnID, call); err != nil { t.Fatal(err) } - if err = s.SubmitFunctionResult(t.Context(), tenant, waiting.ID, waitInput.TurnID, call.CallID, json.RawMessage(`{"success":true,"output":"saved"}`)); err != nil { + if err = SubmitFixtureFunctionResult(t.Context(), s, tenant, waiting.ID, waitInput.TurnID, call.CallID, json.RawMessage(`{"success":true,"output":"saved"}`)); err != nil { t.Fatal(err) } before, err := s.GetSession(t.Context(), tenant, active.ID) @@ -100,8 +125,8 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { _, err = writer.TransitionTurn(t.Context(), tenant, queued.ID, pending.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}) mustReject("claim", err) mustReject("journal", writer.AppendTurnEvents(t.Context(), tenant, active.ID, input.TurnID, 1, []sessions.ExecutionEvent{{Kind: "delta", Payload: json.RawMessage(`{"delta":"stale"}`)}})) - mustReject("callback", writer.RecordFunctionCall(t.Context(), tenant, active.ID, input.TurnID, functionCallFixture("late"))) - mustReject("receipt", writer.ConfirmFunctionResult(t.Context(), tenant, waiting.ID, waitInput.TurnID, call.CallID)) + mustReject("callback", functions.RecordFunctionCall(t.Context(), tenant, active.ID, input.TurnID, functionCallFixture("late"))) + mustReject("receipt", functions.ConfirmFunctionResult(t.Context(), tenant, waiting.ID, waitInput.TurnID, call.CallID)) _, err = writer.CompleteExecution(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnCompleted, json.RawMessage(`{"done":{"content":"stale"}}`), "stale-native", input.Sequence) mustReject("completion", err) _, err = writer.TransitionTurn(t.Context(), tenant, active.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnInProgress, Status: sessions.TurnFailed}) @@ -121,7 +146,7 @@ func TestExecutionLeaseLossFencesAllLifecycleWrites(t *testing.T) { if err != nil || len(events) != 0 { t.Fatal("stale journal persisted", events, err) } - saved, err := s.GetFunctionCall(t.Context(), tenant, waiting.ID, waitInput.TurnID, call.CallID) + saved, err := FixtureFunctionCall(t.Context(), s.pool, tenant, waiting.ID, waitInput.TurnID, call.CallID) if err != nil || saved.Applied { t.Fatal("stale receipt persisted", saved, err) } diff --git a/services/core/internal/store/export_test.go b/services/core/internal/store/export_test.go index 5c1051e26..4dd9d8b7b 100644 --- a/services/core/internal/store/export_test.go +++ b/services/core/internal/store/export_test.go @@ -4,13 +4,19 @@ import ( "bytes" "context" "encoding/json" + "errors" + + "testing" v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/credentialcrypto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/identity" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" - "testing" ) func NewTestStore(t *testing.T) (*Store, *pgxpool.Pool) { return testStore(t) } @@ -81,3 +87,32 @@ func SkillArchive(t *testing.T, marker string) []byte { return skillArchive(t, m func (s *Store) CommitLegacyDeletion(ctx context.Context, tenantID, sessionID string) error { return s.commitLegacyDeletion(ctx, tenantID, sessionID) } + +// FixtureFunctionCall reads a stored function call of the Turn from the test's +// database for assertions. A call outside the tenant's Session and Turn is +// sessions.ErrNotFound. +func FixtureFunctionCall(ctx context.Context, pool *pgxpool.Pool, tenantID, sessionID, turnID, callID string) (sessions.FunctionCall, error) { + p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) + if err != nil { + return sessions.FunctionCall{}, err + } + row, err := sqlc.New(pool).GetFunctionCall(ctx, sqlc.GetFunctionCallParams{TenantID: p.TenantID, SessionID: p.SessionID, TurnID: p.ID, CallID: callID}) + if errors.Is(err, pgx.ErrNoRows) { + return sessions.FunctionCall{}, sessions.ErrNotFound + } + if err != nil { + return sessions.FunctionCall{}, err + } + return sessions.FunctionCall{CallID: row.CallID, ExecutorCallID: row.ExecutorCallID, Name: row.Name, Arguments: row.Arguments, Result: row.Result, Applied: row.Applied}, nil +} + +// SubmitFixtureFunctionResult submits result for the Turn's call as a public +// tool_result input under a fresh request key. +func SubmitFixtureFunctionResult(ctx context.Context, s *Store, tenantID, sessionID, turnID, callID string, result json.RawMessage) error { + payload, err := json.Marshal(sessions.FunctionResultInput{TurnID: turnID, CallID: callID, Result: result}) + if err != nil { + return err + } + _, err = s.SubmitInputs(ctx, tenantID, sessionID, uuid.NewString(), []sessions.Input{{Kind: "tool_result", Payload: payload}}) + return err +} diff --git a/services/core/internal/store/fixture_db_test.go b/services/core/internal/store/fixture_db_test.go index 7ce0713d6..88b605717 100644 --- a/services/core/internal/store/fixture_db_test.go +++ b/services/core/internal/store/fixture_db_test.go @@ -12,6 +12,8 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/deploymentpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/modelconfigurationpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) @@ -59,34 +61,50 @@ func startWorker(t testing.TB, ctx context.Context, db fixtureDB, dispatcher *ex // startWorkerErr is startWorker for tests that assert a startup failure. func startWorkerErr(ctx context.Context, db fixtureDB, dispatcher *execution.Dispatcher) (*execution.Worker, error) { - _, credentials, err := fixtureVaults(db) + lease, err := pgunit.AcquireLease(ctx, db.pool) if err != nil { return nil, err } - sessionStore, sessionService, err := fixtureSessions(db) + owner, err := fixtureOwner(db, dispatcher.Store, lease) if err != nil { - return nil, err + return nil, errors.Join(err, lease.Close(ctx)) } - lease, err := pgunit.AcquireLease(ctx, db.pool) + return startOwnedWorkerErr(ctx, db, dispatcher, owner) +} + +// startOwnedWorker is startWorker on an Owner the test already holds, for tests +// that also run execution operations on it. The Worker closes its lease when +// Run exits. +func startOwnedWorker(t testing.TB, ctx context.Context, db fixtureDB, dispatcher *execution.Dispatcher, owner execution.Owner) *execution.Worker { + t.Helper() + worker, err := startOwnedWorkerErr(ctx, db, dispatcher, owner) if err != nil { - return nil, err + t.Fatal(err) } - deployments, changes, err := fixtureDeploymentExecution(db, lease) + return worker +} + +func startOwnedWorkerErr(ctx context.Context, db fixtureDB, dispatcher *execution.Dispatcher, owner execution.Owner) (*execution.Worker, error) { + _, credentials, err := fixtureVaults(db) if err != nil { - return nil, errors.Join(err, lease.Close(ctx)) + return nil, errors.Join(err, owner.Lease.Close(ctx)) + } + sessionStore, sessionService, err := fixtureSessions(db) + if err != nil { + return nil, errors.Join(err, owner.Lease.Close(ctx)) + } + deployments, err := fixtureDeploymentService(db) + if err != nil { + return nil, errors.Join(err, owner.Lease.Close(ctx)) } owned := *dispatcher owned.Credentials = credentials owned.Observer = modelconfigurationpg.New(pgunit.NewPool(db.pool), db.cipher) owned.Deployment = deployments + owned.DeploymentReader = deploymentpg.New(pgunit.NewPool(db.pool), db.cipher) owned.Sessions = sessionService owned.SessionsReader = sessionStore - owned.DeploymentReader = deploymentpg.New(pgunit.NewPool(db.pool), db.cipher) - return execution.StartWorker(ctx, &owned, execution.Owner{ - Lease: lease, - Store: store.NewExecution(dispatcher.Store, lease), - Deployment: changes, - }) + return execution.StartWorker(ctx, &owned, owner) } // executionOwner acquires the execution lease on db and builds s's execution @@ -99,13 +117,28 @@ func executionOwner(t testing.TB, db fixtureDB, s *store.Store) execution.Owner t.Fatal(err) } t.Cleanup(func() { _ = lease.Close(context.Background()) }) - _, changes, err := fixtureDeploymentExecution(db, lease) + owner, err := fixtureOwner(db, s, lease) if err != nil { t.Fatal(err) } + return owner +} + +// fixtureOwner builds s's execution writer and the deployment and Session +// execution operations on lease, as cmd/server does. +func fixtureOwner(db fixtureDB, s *store.Store, lease *pgunit.Lease) (execution.Owner, error) { + _, changes, err := fixtureDeploymentExecution(db, lease) + if err != nil { + return execution.Owner{}, err + } + sessionExecution, err := sessions.NewExecutionOperations(sessionpg.NewExecution(lease)) + if err != nil { + return execution.Owner{}, err + } return execution.Owner{ Lease: lease, Store: store.NewExecution(s, lease), Deployment: changes, - } + Sessions: sessionExecution, + }, nil } diff --git a/services/core/internal/store/function_calls.go b/services/core/internal/store/function_calls.go deleted file mode 100644 index 95fe59a82..000000000 --- a/services/core/internal/store/function_calls.go +++ /dev/null @@ -1,108 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "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" -) - -// 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 := sessionpg.TurnLookup(tenantID, sessionID, turnID) - if err != nil { - return err - } - if !validFunctionIdentity(call.CallID) || !validFunctionIdentity(call.ExecutorCallID) || !validFunctionIdentity(call.Name) || len(call.Arguments) > 512*1024 || !json.Valid(call.Arguments) || call.Result != nil || call.Applied { - return sessions.ErrInvalidInput - } - return s.withSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - turn, err := q.GetTurn(ctx, p) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - matches, err := q.MatchFunctionCall(ctx, sqlc.MatchFunctionCallParams{SessionID: session, TurnID: p.ID, CallID: call.CallID, ExecutorCallID: call.ExecutorCallID, Name: call.Name, Arguments: call.Arguments}) - if err == nil { - if !matches { - return sessions.ErrIdempotencyConflict - } - return nil - } - if !errors.Is(err, pgx.ErrNoRows) { - return err - } - if !acceptsFunctionResult(turn) { - return sessions.ErrTurnConflict - } - count, err := q.CreateFunctionCall(ctx, sqlc.CreateFunctionCallParams{SessionID: session, TurnID: p.ID, CallID: call.CallID, ExecutorCallID: call.ExecutorCallID, Name: call.Name, Arguments: call.Arguments}) - if err != nil { - return err - } - if count != 1 { - return sessions.ErrIdempotencyConflict - } - return recordFunctionState(ctx, q, turn) - }) -} - -func (s *Store) GetFunctionCall(ctx context.Context, tenantID, sessionID, turnID, callID string) (sessions.FunctionCall, error) { - p, err := sessionpg.TurnLookup(tenantID, sessionID, turnID) - if err != nil { - return sessions.FunctionCall{}, err - } - if !validFunctionIdentity(callID) { - return sessions.FunctionCall{}, sessions.ErrInvalidInput - } - row, err := s.queries.GetFunctionCall(ctx, sqlc.GetFunctionCallParams{TenantID: p.TenantID, SessionID: p.SessionID, TurnID: p.ID, CallID: callID}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.FunctionCall{}, sessions.ErrNotFound - } - if err != nil { - return sessions.FunctionCall{}, err - } - return functionCallFromRow(row), nil -} - -// 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 := sessionpg.TurnLookup(tenantID, sessionID, turnID) - if err != nil { - return nil, err - } - result := make([]sessions.FunctionCall, 0) - err = s.withSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - if _, err := q.GetTurn(ctx, p); errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } else if err != nil { - return err - } - rows, err := q.ListPendingFunctionCalls(ctx, sqlc.ListPendingFunctionCallsParams{SessionID: session, TurnID: p.ID}) - if err != nil { - return err - } - for _, row := range rows { - result = append(result, functionCallFromRow(row)) - } - return nil - }) - return result, err -} - -func functionCallFromRow(row sqlc.FunctionCall) sessions.FunctionCall { - return sessions.FunctionCall{CallID: row.CallID, ExecutorCallID: row.ExecutorCallID, Name: row.Name, Arguments: row.Arguments, Result: row.Result, Applied: row.Applied} -} - -func validFunctionIdentity(id string) bool { return strings.TrimSpace(id) != "" && len(id) <= 512 } - -func acceptsFunctionResult(turn sqlc.Turn) bool { - return (turn.Status == sessions.TurnInProgress || turn.Status == sessions.TurnWaiting) && !turn.CancelRequestedAt.Valid -} diff --git a/services/core/internal/store/function_calls_test.go b/services/core/internal/store/function_calls_test.go deleted file mode 100644 index 446ac5957..000000000 --- a/services/core/internal/store/function_calls_test.go +++ /dev/null @@ -1,271 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "reflect" - "strings" - "sync" - "testing" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/google/uuid" -) - -func functionCallFixture(id string) sessions.FunctionCall { - return sessions.FunctionCall{CallID: id, ExecutorCallID: "native-" + id, Name: "lookup", Arguments: json.RawMessage(`{"ticket":9007199254740993}`)} -} - -func TestFunctionCallsPersistCompleteResultsAndReceipts(t *testing.T) { - s, pool := testStore(t) - tenant, session := newTurnSession(t, s) - turn := submitMessage(t, s, tenant, session.ID, "start").TurnID - transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) - transition(t, s, tenant, session.ID, turn, sessions.TurnInProgress, sessions.TurnWaiting) - results := []string{ - `{"success":true,"output":"answer"}`, - `{"success":true,"output":""}`, - `{"success":false,"output":[{"type":"input_text","text":"before"},{"type":"input_image","image_url":"data:image/png;base64,AA=="},{"type":"input_text","text":""}],"error":"tool failed"}`, - `{"success":true,"output":[],"error":null}`, - `{"success":false,"output":null,"error":"missing"}`, - `{"success":false}`, - } - for i, raw := range results { - call := functionCallFixture(fmt.Sprint(i)) - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, call); err != nil { - t.Fatal(err) - } - retry := call - retry.Arguments = json.RawMessage(`{ "ticket" : 9007199254740993 }`) - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, retry); err != nil { - t.Fatal(err) - } - if err := s.ConfirmFunctionResult(t.Context(), tenant, session.ID, turn, call.CallID); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal("unsubmitted result applied", err) - } - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, call.CallID, json.RawMessage(raw)); err != nil { - t.Fatal(err) - } - } - before, err := s.PendingFunctionCalls(t.Context(), tenant, session.ID, turn) - if err != nil || len(before) != len(results) { - t.Fatal(before, err) - } - pool.Close() - reopened, _ := testStore(t) - restored, err := reopened.PendingFunctionCalls(t.Context(), tenant, session.ID, turn) - if err != nil || !reflect.DeepEqual(before, restored) { - t.Fatal("recovery lost calls/results", restored, err) - } - for i, expected := range results { - id := fmt.Sprint(i) - row, err := reopened.GetFunctionCall(t.Context(), tenant, session.ID, turn, id) - if err != nil || row.Applied || row.ExecutorCallID != "native-"+id || !strings.Contains(string(row.Arguments), "9007199254740993") { - t.Fatal(row, err) - } - normalized, err := jsonobject.Normalize(row.Result) - wanted, _ := jsonobject.Normalize(json.RawMessage(expected)) - if err != nil || string(normalized) != string(wanted) { - t.Fatal("result changed", string(row.Result), err) - } - if err := reopened.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, id, json.RawMessage(expected)); err != nil { - t.Fatal(err) - } - if err := reopened.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, id, json.RawMessage(`{"success":false,"error":"changed"}`)); !errors.Is(err, sessions.ErrFunctionResultConflict) { - t.Fatal(err) - } - for range 2 { - if err := reopened.ConfirmFunctionResult(t.Context(), tenant, session.ID, turn, id); err != nil { - t.Fatal(err) - } - } - } - pending, err := reopened.PendingFunctionCalls(t.Context(), tenant, session.ID, turn) - if err != nil || len(pending) != 0 { - t.Fatal(pending, err) - } - again, _ := testStore(t) - row, err := again.GetFunctionCall(t.Context(), tenant, session.ID, turn, "0") - if err != nil || !row.Applied { - t.Fatal("receipt did not persist", row, err) - } -} - -func TestFunctionCallsAreScopedAndImmutable(t *testing.T) { - s, _ := testStore(t) - tenant, session := newTurnSession(t, s) - turn := submitMessage(t, s, tenant, session.ID, "start").TurnID - call := functionCallFixture("call") - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, call); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal("queued turn accepted callback", err) - } - transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, call); err != nil { - t.Fatal(err) - } - for _, field := range []string{"call", "executor", "name", "arguments"} { - changed := call - switch field { - case "call": - changed.CallID = "another-public-id" - case "executor": - changed.ExecutorCallID = "another-executor-id" - case "name": - changed.Name = "another-function" - case "arguments": - changed.Arguments = json.RawMessage(`{"ticket":9007199254740992}`) - } - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, changed); !errors.Is(err, sessions.ErrIdempotencyConflict) { - t.Fatal(field, err) - } - } - for _, scope := range []struct{ tenant, session, turn string }{ - {uuid.NewString(), session.ID, turn}, {tenant, uuid.NewString(), turn}, {tenant, session.ID, uuid.NewString()}, - } { - _, err := s.GetFunctionCall(t.Context(), scope.tenant, scope.session, scope.turn, "call") - if !errors.Is(err, sessions.ErrNotFound) { - t.Fatal(err) - } - if _, err := s.PendingFunctionCalls(t.Context(), scope.tenant, scope.session, scope.turn); !errors.Is(err, sessions.ErrNotFound) { - t.Fatal(err) - } - if err := s.RecordFunctionCall(t.Context(), scope.tenant, scope.session, scope.turn, call); !errors.Is(err, sessions.ErrNotFound) { - t.Fatal(err) - } - if err := s.SubmitFunctionResult(t.Context(), scope.tenant, scope.session, scope.turn, "call", json.RawMessage(`{"success":true}`)); !errors.Is(err, sessions.ErrNotFound) { - t.Fatal(err) - } - if err := s.ConfirmFunctionResult(t.Context(), scope.tenant, scope.session, scope.turn, "call"); !errors.Is(err, sessions.ErrNotFound) { - t.Fatal(err) - } - } - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, "missing", json.RawMessage(`{"success":true}`)); !errors.Is(err, sessions.ErrNotFound) { - t.Fatal(err) - } -} - -func TestFunctionResultsCannotApplyAfterCancellationOrCompletion(t *testing.T) { - for _, terminal := range []string{sessions.TurnCancelled, sessions.TurnCompleted, sessions.TurnFailed} { - t.Run(terminal, func(t *testing.T) { - s, _ := testStore(t) - tenant, session := newTurnSession(t, s) - turn := submitMessage(t, s, tenant, session.ID, "start").TurnID - transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) - for _, id := range []string{"submitted", "pending"} { - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture(id)); err != nil { - t.Fatal(err) - } - } - result := json.RawMessage(`{"success":true,"output":"saved"}`) - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, "submitted", result); err != nil { - t.Fatal(err) - } - if terminal == sessions.TurnCancelled { - if _, err := s.RequestCancel(t.Context(), tenant, session.ID, "cancel"); err != nil { - t.Fatal(err) - } - assertNoPendingFunctions(t, s, tenant, session.ID, turn) - if err := s.ConfirmFunctionResult(t.Context(), tenant, session.ID, turn, "submitted"); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal(err) - } - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, "pending", result); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal(err) - } - } - transition(t, s, tenant, session.ID, turn, sessions.TurnWaiting, terminal) - assertNoPendingFunctions(t, s, tenant, session.ID, turn) - if err := s.ConfirmFunctionResult(t.Context(), tenant, session.ID, turn, "submitted"); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal(err) - } - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, "pending", result); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal(err) - } - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture("late")); !errors.Is(err, sessions.ErrTurnConflict) { - t.Fatal(err) - } - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, "submitted", result); err != nil { - t.Fatal("identical retry changed its result", err) - } - row, err := s.GetFunctionCall(t.Context(), tenant, session.ID, turn, "submitted") - if err != nil || row.Applied || len(row.Result) == 0 { - t.Fatal("historical result lost or falsely applied", row, err) - } - next := submitMessage(t, s, tenant, session.ID, "next").TurnID - transition(t, s, tenant, session.ID, next, sessions.TurnQueued, sessions.TurnInProgress) - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, next, functionCallFixture("pending")); err != nil { - t.Fatal("call identity leaked across Turns", err) - } - }) - } -} - -func assertNoPendingFunctions(t *testing.T, s *Store, tenant, session, turn string) { - t.Helper() - rows, err := s.PendingFunctionCalls(t.Context(), tenant, session, turn) - if err != nil || len(rows) != 0 { - t.Fatal(rows, err) - } -} - -func TestFunctionResultConcurrentSubmissionsChooseOneValue(t *testing.T) { - s, _ := testStore(t) - other, _ := testStore(t) - tenant, session := newTurnSession(t, s) - turn := submitMessage(t, s, tenant, session.ID, "start").TurnID - transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture("call")); err != nil { - t.Fatal(err) - } - var wg sync.WaitGroup - outcomes := make(chan error, 2) - for _, output := range []string{"first", "second"} { - wg.Add(1) - go func() { - defer wg.Done() - result, _ := json.Marshal(map[string]any{"success": true, "output": output}) - outcomes <- other.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, "call", result) - }() - } - wg.Wait() - close(outcomes) - winners, conflicts := 0, 0 - for err := range outcomes { - if err == nil { - winners++ - } else if errors.Is(err, sessions.ErrFunctionResultConflict) { - conflicts++ - } else { - t.Fatal(err) - } - } - if winners != 1 || conflicts != 1 { - t.Fatal(winners, conflicts) - } -} - -func TestFunctionResultInvalidStorageInputDoesNotConsumeCall(t *testing.T) { - s, _ := testStore(t) - tenant, session := newTurnSession(t, s) - turn := submitMessage(t, s, tenant, session.ID, "start").TurnID - transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture("call")); err != nil { - t.Fatal(err) - } - for _, raw := range []string{"", "null", "[]", "{} {}", `{"output":"` + strings.Repeat("a", 512*1024) + `"}`} { - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, "call", json.RawMessage(raw)); !errors.Is(err, sessions.ErrInvalidInput) { - t.Fatal(err) - } - } - ctx, cancel := context.WithCancel(t.Context()) - cancel() - if err := s.SubmitFunctionResult(ctx, tenant, session.ID, turn, "call", json.RawMessage(`{"success":true}`)); !errors.Is(err, context.Canceled) { - t.Fatal(err) - } - row, err := s.GetFunctionCall(t.Context(), tenant, session.ID, turn, "call") - if err != nil || row.Result != nil || row.Applied { - t.Fatal(row, err) - } -} diff --git a/services/core/internal/store/function_execution_native_test.go b/services/core/internal/store/function_execution_native_test.go index 1e71b44fb..99a60a802 100644 --- a/services/core/internal/store/function_execution_native_test.go +++ b/services/core/internal/store/function_execution_native_test.go @@ -44,12 +44,12 @@ func TestNativeFunctionExecutionPersistsCallsResultsAndContinuity(t *testing.T) } raw, _ := json.Marshal(value) for range 2 { - if err := h.s.SubmitFunctionResult(ctx, h.tenant, h.session.ID, input.TurnID, action.CallID, raw); err != nil { + if err := store.SubmitFixtureFunctionResult(ctx, h.s, h.tenant, h.session.ID, input.TurnID, action.CallID, raw); err != nil { t.Fatal(err) } } h.finished(running, sessions.TurnCompleted) - saved, err := h.s.GetFunctionCall(ctx, h.tenant, h.session.ID, input.TurnID, action.CallID) + saved, err := store.FixtureFunctionCall(ctx, h.db.pool, h.tenant, h.session.ID, input.TurnID, action.CallID) if err != nil || !saved.Applied { t.Fatal(saved, err) } diff --git a/services/core/internal/store/function_execution_test.go b/services/core/internal/store/function_execution_test.go index 199dee6ec..58aa56681 100644 --- a/services/core/internal/store/function_execution_test.go +++ b/services/core/internal/store/function_execution_test.go @@ -81,7 +81,7 @@ func TestExecutionFunctionsWaitForEveryApplicationReceipt(t *testing.T) { } for _, id := range []string{"a", "b"} { public := items.Identity(input.TurnID, "tool:"+id) - if err := h.s.SubmitFunctionResult(t.Context(), h.tenant, h.session.ID, input.TurnID, public, json.RawMessage(`{"success":true,"output":"saved"}`)); err != nil { + if err := store.SubmitFixtureFunctionResult(t.Context(), h.s, h.tenant, h.session.ID, input.TurnID, public, json.RawMessage(`{"success":true,"output":"saved"}`)); err != nil { t.Fatal(err) } } @@ -89,7 +89,7 @@ func TestExecutionFunctionsWaitForEveryApplicationReceipt(t *testing.T) { var reply proto.FunctionResultPayload _ = h.read(proto.TypeFunctionResult).DecodePayload(&reply) public := items.Identity(input.TurnID, "tool:"+reply.CallID) - saved, err := h.s.GetFunctionCall(t.Context(), h.tenant, h.session.ID, input.TurnID, public) + saved, err := store.FixtureFunctionCall(t.Context(), h.db.pool, h.tenant, h.session.ID, input.TurnID, public) if err != nil || saved.Applied || reply.DeliveryID != "function:"+public || len(reply.Content) != 1 || *reply.Content[0].Text != "saved" { t.Fatal(saved, reply, err) } @@ -127,7 +127,7 @@ func TestExecutionFunctionsCancellationAndUnconfirmedResults(t *testing.T) { h.write(input.TurnID, proto.TypeFunctionCall, proto.FunctionCallPayload{CallID: "a", Name: "lookup_ticket", Arguments: json.RawMessage(`{}`)}) state := functionState(t, h, 1) id := state.RequiredActions[0].CallID - if err := h.s.SubmitFunctionResult(t.Context(), h.tenant, h.session.ID, input.TurnID, id, json.RawMessage(`{"success":false,"error":"tool failed"}`)); err != nil { + if err := store.SubmitFixtureFunctionResult(t.Context(), h.s, h.tenant, h.session.ID, input.TurnID, id, json.RawMessage(`{"success":false,"error":"tool failed"}`)); err != nil { t.Fatal(err) } var reply proto.FunctionResultPayload @@ -160,11 +160,11 @@ func TestExecutionFunctionsCancellationAndUnconfirmedResults(t *testing.T) { t.Fatal(bound, err) } } - saved, err := h.s.GetFunctionCall(t.Context(), h.tenant, h.session.ID, input.TurnID, id) + saved, err := store.FixtureFunctionCall(t.Context(), h.db.pool, h.tenant, h.session.ID, input.TurnID, id) if err != nil || saved.Applied || len(saved.Result) == 0 { t.Fatal(saved, err) } - if err := h.s.ConfirmFunctionResult(t.Context(), h.tenant, h.session.ID, input.TurnID, id); !errors.Is(err, sessions.ErrTurnConflict) { + if err := h.owner().Sessions.ConfirmFunctionResult(t.Context(), h.tenant, h.session.ID, input.TurnID, id); !errors.Is(err, sessions.ErrTurnConflict) { t.Fatal(err) } functionState(t, h, 0) diff --git a/services/core/internal/store/function_images_native_test.go b/services/core/internal/store/function_images_native_test.go index 8551062dc..4d660edc6 100644 --- a/services/core/internal/store/function_images_native_test.go +++ b/services/core/internal/store/function_images_native_test.go @@ -11,6 +11,7 @@ import ( "time" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -72,7 +73,7 @@ func TestNativeFunctionImagePublicExecution(t *testing.T) { t.Fatal("invalid evidence", err) } for _, item := range proof.Calls { - call, err := h.s.GetFunctionCall(ctx, h.tenant, proof.Session, item.Turn, item.Call) + call, err := store.FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, item.Turn, item.Call) if err != nil || !call.Applied { t.Fatal("function delivery acknowledgement missing", err) } diff --git a/services/core/internal/store/function_input_execution_test.go b/services/core/internal/store/function_input_execution_test.go index 09444f86d..fd0964274 100644 --- a/services/core/internal/store/function_input_execution_test.go +++ b/services/core/internal/store/function_input_execution_test.go @@ -8,6 +8,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestExecutionFunctionInputBatchStillSteersMessages(t *testing.T) { @@ -52,7 +53,7 @@ func TestExecutionFunctionInputBatchStillSteersMessages(t *testing.T) { } h.write(input.TurnID, proto.TypeDone, proto.DonePayload{Content: "done"}) h.finished(running, sessions.TurnCompleted) - saved, err := h.s.GetFunctionCall(t.Context(), h.tenant, h.session.ID, input.TurnID, state.RequiredActions[0].CallID) + saved, err := store.FixtureFunctionCall(t.Context(), h.db.pool, h.tenant, h.session.ID, input.TurnID, state.RequiredActions[0].CallID) if err != nil || !saved.Applied { t.Fatal(saved, err) } diff --git a/services/core/internal/store/function_inputs.go b/services/core/internal/store/function_inputs.go deleted file mode 100644 index 4da8ad754..000000000 --- a/services/core/internal/store/function_inputs.go +++ /dev/null @@ -1,80 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/jsonobject" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgtype" -) - -func functionInput(raw json.RawMessage) (sessions.FunctionResultInput, error) { - var input sessions.FunctionResultInput - if json.Unmarshal(raw, &input) != nil || input.TurnID == "" || !validFunctionIdentity(input.CallID) || len(input.Result) == 0 { - return input, sessions.ErrInvalidInput - } - result, err := jsonobject.Normalize(input.Result) - if err != nil { - return input, fmt.Errorf("%w: %w", sessions.ErrInvalidInput, err) - } - input.Result = result - return input, nil -} - -func admitFunctionResult(ctx context.Context, q *sqlc.Queries, tenantID string, session pgtype.UUID, key string, position int32, input sessions.Input) (sessions.InputReceipt, error) { - result, err := functionInput(input.Payload) - if err != nil { - return sessions.InputReceipt{}, err - } - tenant, err := parseID(tenantID) - if err != nil { - return sessions.InputReceipt{}, err - } - var turn sqlc.Turn - found := false - // A malformed Turn ID cannot name a Turn; it is classified like an unknown one. - if id, err := parseID(result.TurnID); err == nil { - turn, err = q.GetTurn(ctx, sqlc.GetTurnParams{TenantID: tenant, SessionID: session, ID: id}) - if err == nil { - found = true - } else if !errors.Is(err, pgx.ErrNoRows) { - return sessions.InputReceipt{}, err - } - } - if found { - err = storeFunctionResult(ctx, q, turn, result.CallID, result.Result) - if errors.Is(err, sessions.ErrNotFound) { - found = false - } else if err != nil { - return sessions.InputReceipt{}, err - } - } - if !found { - return sessions.InputReceipt{}, unknownFunctionResultTarget(ctx, q, session, result.CallID) - } - sequence, err := q.CreateTurnInput(ctx, sqlc.CreateTurnInputParams{ - SessionID: session, TurnID: turn.ID, IdempotencyKey: key, Kind: input.Kind, Payload: input.Payload, BatchPosition: position, - }) - if err != nil { - return sessions.InputReceipt{}, err - } - return inputReceipt(sequence, turn.ID, false), nil -} - -// unknownFunctionResultTarget classifies a result whose Turn does not own the -// call. The answer depends only on this Session's own calls. -func unknownFunctionResultTarget(ctx context.Context, q *sqlc.Queries, session pgtype.UUID, callID string) error { - elsewhere, err := q.SessionHasFunctionCall(ctx, sqlc.SessionHasFunctionCallParams{SessionID: session, CallID: callID}) - if err != nil { - return err - } - if elsewhere { - return sessions.ErrFunctionCallTurnMismatch - } - return sessions.ErrUnknownFunctionCall -} diff --git a/services/core/internal/store/function_inputs_public_test.go b/services/core/internal/store/function_inputs_public_test.go index 005acb0c9..72614c5d8 100644 --- a/services/core/internal/store/function_inputs_public_test.go +++ b/services/core/internal/store/function_inputs_public_test.go @@ -41,8 +41,9 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { if _, err := s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } + functions := executionOwner(t, db, s).Sessions for _, id := range []string{"a", "b", "c", "rollback", "late"} { - if err := s.RecordFunctionCall(ctx, tenant, session.ID, input.TurnID, sessions.FunctionCall{CallID: id, ExecutorCallID: "native-" + id, Name: "lookup", Arguments: json.RawMessage(`{}`)}); err != nil { + if err := functions.RecordFunctionCall(ctx, tenant, session.ID, input.TurnID, sessions.FunctionCall{CallID: id, ExecutorCallID: "native-" + id, Name: "lookup", Arguments: json.RawMessage(`{}`)}); err != nil { t.Fatal(err) } } @@ -72,7 +73,7 @@ func TestFunctionInputsOfficialClientAtomicAdmission(t *testing.T) { t.Fatalf("SDK admission: %s %v %s", line, err, stderr.String()) } for _, id := range []string{"a", "b", "c", "rollback", "late"} { - call, err := s.GetFunctionCall(ctx, tenant, session.ID, input.TurnID, id) + call, err := store.FixtureFunctionCall(ctx, db.pool, tenant, session.ID, input.TurnID, id) if err != nil || call.Applied { t.Fatal(call, err) } diff --git a/services/core/internal/store/function_inputs_test.go b/services/core/internal/store/function_inputs_test.go index bdb6b4179..53708dc1c 100644 --- a/services/core/internal/store/function_inputs_test.go +++ b/services/core/internal/store/function_inputs_test.go @@ -22,13 +22,13 @@ func resultInput(t *testing.T, turn, call, result string) sessions.Input { return sessions.Input{Kind: "tool_result", Payload: raw} } -func functionInputFixture(t *testing.T, s *Store) (string, sessions.Session, string) { +func functionInputFixture(t *testing.T, s *Store, functions *sessions.ExecutionOperations) (string, sessions.Session, string) { t.Helper() tenant, session := newTurnSession(t, s) turn := submitMessage(t, s, tenant, session.ID, "start").TurnID transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) for _, id := range []string{"a", "b"} { - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture(id)); err != nil { + if err := functions.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture(id)); err != nil { t.Fatal(err) } } @@ -37,7 +37,7 @@ func functionInputFixture(t *testing.T, s *Store) (string, sessions.Session, str func TestFunctionInputBatchesPersistAndReplayWithoutRetargeting(t *testing.T) { s, pool := testStore(t) - tenant, session, turn := functionInputFixture(t, s) + tenant, session, turn := functionInputFixture(t, s, functionExecution(t)) full := `{"success":false,"output":[{"type":"input_text","text":""},{"type":"input_image","image_url":"data:image/png;base64,AA=="},{"type":"input_text","text":"after"}],"error":"failed"}` batch := []sessions.Input{resultInput(t, turn, "a", full), {Kind: "message", Payload: json.RawMessage(`{"text":"Follow up"}`)}, resultInput(t, turn, "b", `{"success":true,"output":null,"error":null}`), {Kind: "cancel", Payload: json.RawMessage(`{}`)}} receipts, err := s.SubmitInputs(t.Context(), tenant, session.ID, "batch", batch) @@ -49,7 +49,7 @@ func TestFunctionInputBatchesPersistAndReplayWithoutRetargeting(t *testing.T) { t.Fatal(receipts) } } - call, err := s.GetFunctionCall(t.Context(), tenant, session.ID, turn, "a") + call, err := FixtureFunctionCall(t.Context(), s.pool, tenant, session.ID, turn, "a") got, _ := jsonobject.Normalize(call.Result) want, _ := jsonobject.Normalize(json.RawMessage(full)) if err != nil || call.Applied || string(got) != string(want) { @@ -93,7 +93,7 @@ func TestFunctionInputBatchesPersistAndReplayWithoutRetargeting(t *testing.T) { if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "same-result", batch[:1]); err != nil { t.Fatal(err) } - call, err = s.GetFunctionCall(t.Context(), tenant, session.ID, turn, "a") + call, err = FixtureFunctionCall(t.Context(), s.pool, tenant, session.ID, turn, "a") if err != nil || call.Applied { t.Fatal(call, err) } @@ -103,7 +103,8 @@ func TestFunctionInputBatchFailureRollsBackEveryWrite(t *testing.T) { for _, mode := range []string{"missing-call", "foreign-turn", "same-tenant-turn", "cancel-first", "changed-result"} { t.Run(mode, func(t *testing.T) { s, _ := testStore(t) - tenant, session, turn := functionInputFixture(t, s) + functions := functionExecution(t) + tenant, session, turn := functionInputFixture(t, s, functions) message := sessions.Input{Kind: "message", Payload: json.RawMessage(`{"text":"Must roll back"}`)} cancel := sessions.Input{Kind: "cancel", Payload: json.RawMessage(`{}`)} first := resultInput(t, turn, "a", `{"success":true}`) @@ -123,7 +124,7 @@ func TestFunctionInputBatchFailureRollsBackEveryWrite(t *testing.T) { } otherTurn := submitMessage(t, s, otherTenant, other.ID, "start").TurnID transition(t, s, otherTenant, other.ID, otherTurn, sessions.TurnQueued, sessions.TurnInProgress) - if err := s.RecordFunctionCall(t.Context(), otherTenant, other.ID, otherTurn, functionCallFixture("a")); err != nil { + if err := functions.RecordFunctionCall(t.Context(), otherTenant, other.ID, otherTurn, functionCallFixture("a")); err != nil { t.Fatal(err) } batch = append(batch, resultInput(t, otherTurn, "a", `{"success":true}`)) @@ -138,7 +139,7 @@ func TestFunctionInputBatchFailureRollsBackEveryWrite(t *testing.T) { if _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "failed-batch", batch); !errors.Is(err, expected) { t.Fatal(err) } - call, err := s.GetFunctionCall(t.Context(), tenant, session.ID, turn, "a") + call, err := FixtureFunctionCall(t.Context(), s.pool, tenant, session.ID, turn, "a") if err != nil || call.Result != nil || call.Applied { t.Fatal(call, err) } @@ -160,7 +161,7 @@ func TestFunctionInputBatchFailureRollsBackEveryWrite(t *testing.T) { func TestFunctionInputConcurrentBatchesSelectOneResult(t *testing.T) { s, _ := testStore(t) other, _ := testStore(t) - tenant, session, turn := functionInputFixture(t, s) + tenant, session, turn := functionInputFixture(t, s, functionExecution(t)) var wg sync.WaitGroup results := make(chan error, 2) for i := range 2 { @@ -195,7 +196,7 @@ func TestFunctionInputConcurrentBatchesSelectOneResult(t *testing.T) { func TestFunctionInputsRejectInvalidTargetsAndStorageObjects(t *testing.T) { s, _ := testStore(t) - tenant, session, turn := functionInputFixture(t, s) + tenant, session, turn := functionInputFixture(t, s, functionExecution(t)) for _, raw := range []string{`{}`, `{"turn_id":"","call_id":"a","result":{}}`, fmt.Sprintf(`{"turn_id":%q,"call_id":" ","result":{}}`, turn), fmt.Sprintf(`{"turn_id":%q,"call_id":"a"}`, turn), fmt.Sprintf(`{"turn_id":%q,"call_id":"a","result":null}`, turn), fmt.Sprintf(`{"turn_id":%q,"call_id":"a","result":[]}`, turn)} { _, err := s.SubmitInputs(t.Context(), tenant, session.ID, "invalid", []sessions.Input{{Kind: "tool_result", Payload: json.RawMessage(raw)}}) if !errors.Is(err, sessions.ErrInvalidInput) { diff --git a/services/core/internal/store/function_item_events_test.go b/services/core/internal/store/function_item_events_test.go index 7d8d63597..7cb36c70d 100644 --- a/services/core/internal/store/function_item_events_test.go +++ b/services/core/internal/store/function_item_events_test.go @@ -64,14 +64,15 @@ func TestFunctionResultItemsRetainSubmittedFields(t *testing.T) { } { t.Run(raw, func(t *testing.T) { s, pool := testStore(t) + functions := functionExecution(t) tenant, session := newTurnSession(t, s) turn := submitMessage(t, s, tenant, session.ID, "start").TurnID transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) call := functionCallFixture(items.Identity(turn, "tool:call")) - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, call); err != nil { + if err := functions.RecordFunctionCall(t.Context(), tenant, session.ID, turn, call); err != nil { t.Fatal(err) } - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, call.CallID, json.RawMessage(raw)); err != nil { + if err := SubmitFixtureFunctionResult(t.Context(), s, tenant, session.ID, turn, call.CallID, json.RawMessage(raw)); err != nil { t.Fatal(err) } event := sessions.ExecutionEvent{Kind: "tool_call", Payload: json.RawMessage(`{"id":"call","stage":"after","observation":{"status":"completed","kind":"function","name":"lookup","arguments":{},"content":[{"type":"input_text","text":"normalized"}]}}`)} diff --git a/services/core/internal/store/function_public_native_test.go b/services/core/internal/store/function_public_native_test.go index 369a14313..c96d5a1f8 100644 --- a/services/core/internal/store/function_public_native_test.go +++ b/services/core/internal/store/function_public_native_test.go @@ -12,6 +12,7 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -43,7 +44,7 @@ func TestNativePublicFunctionExecution(t *testing.T) { t.Fatal(proof, err) } for i, callID := range proof.Calls { - call, err := h.s.GetFunctionCall(ctx, h.tenant, proof.Session, proof.Turns[i], callID) + call, err := store.FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turns[i], callID) if err != nil || call.Applied != (i < 2) { t.Fatal(call, err) } diff --git a/services/core/internal/store/function_results.go b/services/core/internal/store/function_results.go deleted file mode 100644 index ebe4ace4f..000000000 --- a/services/core/internal/store/function_results.go +++ /dev/null @@ -1,92 +0,0 @@ -package store - -import ( - "context" - "encoding/json" - "errors" - "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" - "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/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" -) - -// SubmitFunctionResult stores a caller-validated result object; its wire schema belongs to the API. -func (s *Store) SubmitFunctionResult(ctx context.Context, tenantID, sessionID, turnID, callID string, result json.RawMessage) error { - if len(result) == 0 || len(result) > 512*1024 { - return sessions.ErrInvalidInput - } - result, err := jsonobject.Normalize(result) - if err != nil { - return fmt.Errorf("%w: %w", sessions.ErrInvalidInput, err) - } - return s.withFunctionCall(ctx, tenantID, sessionID, turnID, callID, func(ctx context.Context, q *sqlc.Queries, turn sqlc.Turn, call sqlc.FunctionCall) error { - return storeFunctionResult(ctx, q, turn, call.CallID, result) - }) -} - -// ConfirmFunctionResult records native application, not external tool success. -func (s *Store) ConfirmFunctionResult(ctx context.Context, tenantID, sessionID, turnID, callID string) error { - return s.withFunctionCall(ctx, tenantID, sessionID, turnID, callID, func(ctx context.Context, q *sqlc.Queries, turn sqlc.Turn, call sqlc.FunctionCall) error { - if call.Applied { - return nil - } - if len(call.Result) == 0 || !acceptsFunctionResult(turn) { - return sessions.ErrTurnConflict - } - if err := q.ApplyFunctionResult(ctx, sqlc.ApplyFunctionResultParams{SessionID: turn.SessionID, TurnID: turn.ID, CallID: call.CallID}); err != nil { - return err - } - return recordFunctionState(ctx, q, turn) - }) -} - -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 := sessionpg.TurnLookup(tenantID, sessionID, turnID) - if err != nil { - return err - } - if !validFunctionIdentity(callID) { - return sessions.ErrInvalidInput - } - return s.withSession(ctx, tenantID, sessionID, func(ctx context.Context, q *sqlc.Queries, session pgtype.UUID) error { - turn, err := q.GetTurn(ctx, p) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - call, err := q.GetFunctionCall(ctx, sqlc.GetFunctionCallParams{TenantID: p.TenantID, SessionID: session, TurnID: p.ID, CallID: callID}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - return fn(ctx, q, turn, call) - }) -} - -func storeFunctionResult(ctx context.Context, q *sqlc.Queries, turn sqlc.Turn, callID string, result json.RawMessage) error { - match, err := q.MatchFunctionResult(ctx, sqlc.MatchFunctionResultParams{SessionID: turn.SessionID, TurnID: turn.ID, CallID: callID, Result: result}) - if errors.Is(err, pgx.ErrNoRows) { - return sessions.ErrNotFound - } - if err != nil { - return err - } - if match.Submitted { - if !match.Matches { - return sessions.ErrFunctionResultConflict - } - return nil - } - if !acceptsFunctionResult(turn) { - return sessions.ErrTurnConflict - } - return q.SubmitFunctionResult(ctx, sqlc.SubmitFunctionResultParams{SessionID: turn.SessionID, TurnID: turn.ID, CallID: callID, Result: result}) -} diff --git a/services/core/internal/store/function_state.go b/services/core/internal/store/function_state.go deleted file mode 100644 index c557deefd..000000000 --- a/services/core/internal/store/function_state.go +++ /dev/null @@ -1,52 +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/google/uuid" -) - -func functionActions(ctx context.Context, q *sqlc.Queries, turn sqlc.Turn) ([]v1.FunctionCallAction, error) { - actions := make([]v1.FunctionCallAction, 0) - if !acceptsFunctionResult(turn) { - return actions, nil - } - calls, err := q.ListPendingFunctionCalls(ctx, sqlc.ListPendingFunctionCallsParams{SessionID: turn.SessionID, TurnID: turn.ID}) - if err != nil { - return nil, err - } - for _, call := range calls { - actions = append(actions, v1.FunctionCallAction{Type: "function_call", CallID: call.CallID, Name: call.Name, TurnID: uuid.UUID(turn.ID.Bytes).String(), Arguments: json.RawMessage(call.Arguments)}) - } - return actions, nil -} - -func recordFunctionState(ctx context.Context, q *sqlc.Queries, turn sqlc.Turn) error { - actions, err := functionActions(ctx, q, turn) - if err != nil { - return err - } - status := sessions.TurnInProgress - if len(actions) > 0 { - status = sessions.TurnWaiting - } - if turn.Status != status { - turn, err = q.TransitionTurn(ctx, sqlc.TransitionTurnParams{ID: turn.ID, SessionID: turn.SessionID, ExpectedStatus: turn.Status, NewStatus: status, Outcome: []byte(`{}`)}) - if err != nil { - return err - } - if err := sessionpg.AppendChanges(ctx, q, turn.SessionID, sessions.TurnChanges(sessionpg.TurnFromRow(turn), false)...); err != nil { - return err - } - } - usage, err := sessionpg.LoadUsage(ctx, q, turn.SessionID) - if err != nil { - return err - } - return sessionpg.AppendChanges(ctx, q, turn.SessionID, sessions.ActivityChange(sessionpg.TurnFromRow(turn), usage, actions)) -} diff --git a/services/core/internal/store/function_state_public_test.go b/services/core/internal/store/function_state_public_test.go index af5a21325..96b8f8031 100644 --- a/services/core/internal/store/function_state_public_test.go +++ b/services/core/internal/store/function_state_public_test.go @@ -38,9 +38,10 @@ func TestFunctionStateOfficialClientReadsAndLiveEvents(t *testing.T) { if _, err := s.TransitionTurn(ctx, tenant, session.ID, input.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } + functions := executionOwner(t, db, s).Sessions record := func(id string) { t.Helper() - if err := s.RecordFunctionCall(ctx, tenant, session.ID, input.TurnID, sessions.FunctionCall{CallID: id, ExecutorCallID: "private-" + id, Name: "lookup", Arguments: json.RawMessage(`{"ticket":9007199254740993}`)}); err != nil { + if err := functions.RecordFunctionCall(ctx, tenant, session.ID, input.TurnID, sessions.FunctionCall{CallID: id, ExecutorCallID: "private-" + id, Name: "lookup", Arguments: json.RawMessage(`{"ticket":9007199254740993}`)}); err != nil { t.Fatal(err) } } @@ -71,10 +72,10 @@ func TestFunctionStateOfficialClientReadsAndLiveEvents(t *testing.T) { } record("second") for _, id := range []string{"first", "second"} { - if err := s.SubmitFunctionResult(ctx, tenant, session.ID, input.TurnID, id, json.RawMessage(`{"success":true,"output":"private"}`)); err != nil { + if err := store.SubmitFixtureFunctionResult(ctx, s, tenant, session.ID, input.TurnID, id, json.RawMessage(`{"success":true,"output":"private"}`)); err != nil { t.Fatal(err) } - if err := s.ConfirmFunctionResult(ctx, tenant, session.ID, input.TurnID, id); err != nil { + if err := functions.ConfirmFunctionResult(ctx, tenant, session.ID, input.TurnID, id); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/store/function_state_test.go b/services/core/internal/store/function_state_test.go index 9cb4711a4..70dd61af6 100644 --- a/services/core/internal/store/function_state_test.go +++ b/services/core/internal/store/function_state_test.go @@ -13,6 +13,7 @@ import ( func TestFunctionStateSnapshotsRecoveryAndRetries(t *testing.T) { s, pool := testStore(t) + functions := functionExecution(t) tenant, session := newTurnSession(t, s) turn := submitMessage(t, s, tenant, session.ID, "start").TurnID transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) @@ -22,14 +23,14 @@ func TestFunctionStateSnapshotsRecoveryAndRetries(t *testing.T) { } for _, id := range []string{"first", "second"} { for range 2 { - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture(id)); err != nil { + if err := functions.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture(id)); err != nil { t.Fatal(err) } } } assertFunctionState(t, s, tenant, session.ID, sessions.TurnWaiting, 2) for _, id := range []string{"first", "second"} { - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, id, json.RawMessage(`{"success":true,"output":"private result"}`)); err != nil { + if err := SubmitFixtureFunctionResult(t.Context(), s, tenant, session.ID, turn, id, json.RawMessage(`{"success":true,"output":"private result"}`)); err != nil { t.Fatal(err) } } @@ -42,7 +43,7 @@ func TestFunctionStateSnapshotsRecoveryAndRetries(t *testing.T) { } for i, id := range []string{"first", "second"} { for range 2 { - if err := s.ConfirmFunctionResult(t.Context(), tenant, session.ID, turn, id); err != nil { + if err := functions.ConfirmFunctionResult(t.Context(), tenant, session.ID, turn, id); err != nil { t.Fatal(err) } } @@ -87,10 +88,11 @@ func TestFunctionStateCancellationAndTerminalCleanup(t *testing.T) { for _, status := range []string{sessions.TurnCancelled, sessions.TurnFailed, sessions.TurnCompleted} { t.Run(status, func(t *testing.T) { s, _ := testStore(t) + functions := functionExecution(t) tenant, session := newTurnSession(t, s) input := submitMessage(t, s, tenant, session.ID, "start") transition(t, s, tenant, session.ID, input.TurnID, sessions.TurnQueued, sessions.TurnInProgress) - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, input.TurnID, functionCallFixture("call")); err != nil { + if err := functions.RecordFunctionCall(t.Context(), tenant, session.ID, input.TurnID, functionCallFixture("call")); err != nil { t.Fatal(err) } if status == sessions.TurnCancelled { @@ -112,7 +114,7 @@ func TestFunctionStateCancellationAndTerminalCleanup(t *testing.T) { t.Fatal(err) } assertFunctionState(t, s, tenant, session.ID, status, 0) - if _, err := s.GetFunctionCall(t.Context(), tenant, session.ID, input.TurnID, "call"); err != nil { + if _, err := FixtureFunctionCall(t.Context(), s.pool, tenant, session.ID, input.TurnID, "call"); err != nil { t.Fatal("history lost", err) } next := submitMessage(t, s, tenant, session.ID, "next") @@ -126,6 +128,7 @@ func TestFunctionStateCancellationAndTerminalCleanup(t *testing.T) { func TestFunctionStateReadsRemainConsistentDuringReceipts(t *testing.T) { s, _ := testStore(t) + functions := functionExecution(t) tenant, session := newTurnSession(t, s) turn := submitMessage(t, s, tenant, session.ID, "start").TurnID transition(t, s, tenant, session.ID, turn, sessions.TurnQueued, sessions.TurnInProgress) @@ -135,15 +138,15 @@ func TestFunctionStateReadsRemainConsistentDuringReceipts(t *testing.T) { defer close(done) for i := range 30 { id := fmt.Sprint(i) - if err := s.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture(id)); err != nil { + if err := functions.RecordFunctionCall(t.Context(), tenant, session.ID, turn, functionCallFixture(id)); err != nil { t.Error(err) return } - if err := s.SubmitFunctionResult(t.Context(), tenant, session.ID, turn, id, json.RawMessage(`{"success":true}`)); err != nil { + if err := SubmitFixtureFunctionResult(t.Context(), s, tenant, session.ID, turn, id, json.RawMessage(`{"success":true}`)); err != nil { t.Error(err) return } - if err := s.ConfirmFunctionResult(t.Context(), tenant, session.ID, turn, id); err != nil { + if err := functions.ConfirmFunctionResult(t.Context(), tenant, session.ID, turn, id); err != nil { t.Error(err) return } diff --git a/services/core/internal/store/function_stream_native_test.go b/services/core/internal/store/function_stream_native_test.go index 806f54916..9ddf2316e 100644 --- a/services/core/internal/store/function_stream_native_test.go +++ b/services/core/internal/store/function_stream_native_test.go @@ -6,6 +6,8 @@ import ( "os/exec" "path/filepath" "testing" + + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" ) func TestNativePublicFunctionStreamHelper(t *testing.T) { @@ -37,7 +39,7 @@ func TestNativePublicFunctionStreamHelper(t *testing.T) { t.Fatal(proof, err) } for i, callID := range proof.Calls { - call, err := h.s.GetFunctionCall(ctx, h.tenant, proof.Session, proof.Turns[i], callID) + call, err := store.FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turns[i], callID) if err != nil || !call.Applied { t.Fatal(call, err) } diff --git a/services/core/internal/store/input_conflicts_public_test.go b/services/core/internal/store/input_conflicts_public_test.go index a02953e37..fd1343bfb 100644 --- a/services/core/internal/store/input_conflicts_public_test.go +++ b/services/core/internal/store/input_conflicts_public_test.go @@ -52,6 +52,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { server := httptest.NewServer(h) defer server.Close() client := pathIDClient{t: t, server: server} + functions := executionOwner(t, db, s).Sessions create := func(environment string, initial bool) string { t.Helper() @@ -76,14 +77,14 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { if _, err := s.TransitionTurn(ctx, tenant, session, receipt.TurnID, sessions.TurnTransition{ExpectedStatus: sessions.TurnQueued, Status: sessions.TurnInProgress}); err != nil { t.Fatal(err) } - if err := s.RecordFunctionCall(ctx, tenant, session, receipt.TurnID, sessions.FunctionCall{CallID: call, ExecutorCallID: "native-" + call, Name: "lookup", Arguments: json.RawMessage(`{}`)}); err != nil { + if err := functions.RecordFunctionCall(ctx, tenant, session, receipt.TurnID, sessions.FunctionCall{CallID: call, ExecutorCallID: "native-" + call, Name: "lookup", Arguments: json.RawMessage(`{}`)}); err != nil { t.Fatal(err) } return receipt } complete := func(session string, receipt sessions.InputReceipt, call string) { t.Helper() - if err := s.ConfirmFunctionResult(ctx, tenant, session, receipt.TurnID, call); err != nil { + if err := functions.ConfirmFunctionResult(ctx, tenant, session, receipt.TurnID, call); err != nil { t.Fatal(err) } if _, err := s.CompleteExecution(ctx, tenant, session, receipt.TurnID, sessions.TurnCompleted, nil, "", receipt.Sequence); err != nil { @@ -160,7 +161,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { none := `{"type":"none"}` session := create(none, false) first := waiting(session, "first", "first-call") - if err := s.SubmitFunctionResult(ctx, tenant, session, first.TurnID, "first-call", json.RawMessage(`{"success":true,"output":"one"}`)); err != nil { + if err := store.SubmitFixtureFunctionResult(ctx, s, tenant, session, first.TurnID, "first-call", json.RawMessage(`{"success":true,"output":"one"}`)); err != nil { t.Fatal(err) } complete(session, first, "first-call") @@ -232,7 +233,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { if status, body := submit(owner, session, "same-after-completion", valid); status != http.StatusAccepted || body != "" { t.Fatalf("identical result after completion: %d %s", status, body) } - if call, err := s.GetFunctionCall(ctx, tenant, session, current.TurnID, "pending-call"); err != nil || !bytes.Contains(call.Result, []byte(`"value"`)) { + if call, err := store.FixtureFunctionCall(ctx, db.pool, tenant, session, current.TurnID, "pending-call"); err != nil || !bytes.Contains(call.Result, []byte(`"value"`)) { t.Fatal("saved result changed", call, err) } @@ -246,7 +247,7 @@ func TestSessionInputConflictsAndResultTargetsPostgres(t *testing.T) { {"after-cancel", owner, session, []string{late}, 409, turnConflictBody, ""}, {"after-cancel-foreign", foreign, session, []string{late}, 404, missingSessionBody, ""}, }, session) - if call, err := s.GetFunctionCall(ctx, tenant, session, cancelled.TurnID, "late-call"); err != nil || call.Result != nil { + if call, err := store.FixtureFunctionCall(ctx, db.pool, tenant, session, cancelled.TurnID, "late-call"); err != nil || call.Result != nil { t.Fatal("late result was saved", call, err) } diff --git a/services/core/internal/store/message_images_native_test.go b/services/core/internal/store/message_images_native_test.go index bbe9f1aa5..f59e3b23b 100644 --- a/services/core/internal/store/message_images_native_test.go +++ b/services/core/internal/store/message_images_native_test.go @@ -12,6 +12,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -73,7 +74,7 @@ func TestNativeMessageImagePublicExecution(t *testing.T) { if err != nil || json.Unmarshal(raw, &proof) != nil { t.Fatal("invalid evidence", err) } - call, err := h.s.GetFunctionCall(ctx, h.tenant, proof.Session, proof.Turn, proof.Call) + call, err := store.FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turn, proof.Call) if err != nil || !call.Applied { t.Fatal("function application receipt missing", err) } diff --git a/services/core/internal/store/model_protocol_native_test.go b/services/core/internal/store/model_protocol_native_test.go index b1a9c2c88..c46ab363d 100644 --- a/services/core/internal/store/model_protocol_native_test.go +++ b/services/core/internal/store/model_protocol_native_test.go @@ -14,6 +14,7 @@ import ( v1 "github.com/MiniMax-AI/OpenAgentCore/contracts/agents-api/v1" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/modelconfiguration" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -120,7 +121,7 @@ func TestNativeModelProtocolPublicExecution(t *testing.T) { } failed := 0 for _, item := range proof.Calls { - call, err := h.s.GetFunctionCall(ctx, h.tenant, proof.Session, item.Turn, item.Call) + call, err := store.FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, item.Turn, item.Call) if err != nil || !call.Applied { t.Fatal("public function result lacks native delivery acknowledgement") } diff --git a/services/core/internal/store/no_environment_test.go b/services/core/internal/store/no_environment_test.go index b58049d9e..c81e554df 100644 --- a/services/core/internal/store/no_environment_test.go +++ b/services/core/internal/store/no_environment_test.go @@ -20,7 +20,7 @@ func TestNoEnvironmentRejectsUnadvertisedDeviceBeforeClaim(t *testing.T) { t.Fatal(err) } input := h.message("first", "Answer") - if _, err = h.d.Run(ctx, h.tenant, h.session.ID, input.TurnID); err == nil { + if _, err = h.bound().Run(ctx, h.tenant, h.session.ID, input.TurnID); err == nil { t.Fatal("unsupported environment admitted") } turn, err := h.s.GetTurn(ctx, h.tenant, h.session.ID, input.TurnID) diff --git a/services/core/internal/store/prepared_dispatch_test.go b/services/core/internal/store/prepared_dispatch_test.go index 0e0b31158..a05dac21e 100644 --- a/services/core/internal/store/prepared_dispatch_test.go +++ b/services/core/internal/store/prepared_dispatch_test.go @@ -23,8 +23,7 @@ func preparedDispatchHarness(t *testing.T) (*dispatchHarness, sessions.Environme t.Helper() h := newDispatchHarnessForSession(t, []byte(`{"agent":{"model":"test-model","instructions":"Keep this instruction."},"environment":{"type":"self_hosted","workspace_directory":"/workspace"}}`), true) assertNoRuntimeAllocation(t, h) - owner := executionOwner(t, h.db, h.s) - h.d.Store, h.lease = owner.Store, owner.Lease + h.d, h.lease = h.bound(), h.owner().Lease enableWorkerEnvironment(t, h) pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}, {Kind: "message", Payload: json.RawMessage(`{"text":"second"}`)}}) if err != nil { diff --git a/services/core/internal/store/runtime_worker_recovery_test.go b/services/core/internal/store/runtime_worker_recovery_test.go index 8b6030a0d..0edabc0b6 100644 --- a/services/core/internal/store/runtime_worker_recovery_test.go +++ b/services/core/internal/store/runtime_worker_recovery_test.go @@ -35,8 +35,7 @@ func runtimeWorkerHarness(t *testing.T) (*dispatchHarness, *pgxpool.Pool) { func TestPreparedDispatchKeepsPendingReservationAfterComputeConflict(t *testing.T) { h, _ := runtimeWorkerHarness(t) - owner := executionOwner(t, h.db, h.s) - h.d.Store, h.lease = owner.Store, owner.Lease + h.d, h.lease = h.bound(), h.owner().Lease pending, err := h.s.ReserveEnvironmentInput(t.Context(), h.tenant, h.session.ID, "pending", []sessions.Input{{Kind: "message", Payload: json.RawMessage(`{"text":"first"}`)}}) if err != nil { t.Fatal(err) diff --git a/services/core/internal/store/scheduling.go b/services/core/internal/store/scheduling.go index 032aba557..0cb7a446b 100644 --- a/services/core/internal/store/scheduling.go +++ b/services/core/internal/store/scheduling.go @@ -159,7 +159,7 @@ func readSessionActivity(ctx context.Context, q *sqlc.Queries, session sessions. } turn := sessionpg.TurnFromRow(row) session.LastTurn = &turn - session.RequiredActions, err = functionActions(ctx, q, row) + session.RequiredActions, err = sessionpg.LoadRequiredActions(ctx, q, id, turn) if err != nil { return session, err } diff --git a/services/core/internal/store/session_deletion_execution_test.go b/services/core/internal/store/session_deletion_execution_test.go index 6dbdc1c03..83044b453 100644 --- a/services/core/internal/store/session_deletion_execution_test.go +++ b/services/core/internal/store/session_deletion_execution_test.go @@ -15,7 +15,7 @@ func TestDeletedSessionWaitingTurnSettlesWithoutStoppingWorker(t *testing.T) { h := newFunctionHarness(t) ctx, cancel := context.WithCancel(t.Context()) defer cancel() - worker := startWorker(t, ctx, h.db, h.d) + worker := startOwnedWorker(t, ctx, h.db, h.d, h.owner()) done := make(chan error, 1) go func() { done <- worker.Run(ctx) }() defer func() { @@ -54,7 +54,7 @@ func TestDeletedSessionWaitingTurnSettlesWithoutStoppingWorker(t *testing.T) { if _, err := h.s.GetSession(ctx, h.tenant, h.session.ID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal(err) } - if _, err := h.d.Run(ctx, h.tenant, h.session.ID, input.TurnID); !errors.Is(err, sessions.ErrNotFound) { + if _, err := h.bound().Run(ctx, h.tenant, h.session.ID, input.TurnID); !errors.Is(err, sessions.ErrNotFound) { t.Fatal("deleted Session must be unavailable to new dispatch", err) } h.session = publicSession(t, h, "unrelated") diff --git a/services/core/internal/store/session_deletion_lifecycle_public_test.go b/services/core/internal/store/session_deletion_lifecycle_public_test.go index 34b01200d..b5368fbac 100644 --- a/services/core/internal/store/session_deletion_lifecycle_public_test.go +++ b/services/core/internal/store/session_deletion_lifecycle_public_test.go @@ -39,7 +39,8 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { server := httptest.NewServer(h) defer server.Close() client := pathIDClient{t: t, server: server} - writer := executionOwner(t, db, s).Store + leased := executionOwner(t, db, s) + writer := leased.Store create := func(environment string, initial bool) sessions.Session { t.Helper() @@ -70,7 +71,7 @@ func TestSessionDeletionLifecyclePostgres(t *testing.T) { case "cancel": _, err = s.RequestCancel(ctx, tenant, session.ID, "cancel") case "function": - err = s.RecordFunctionCall(ctx, tenant, session.ID, receipt.TurnID, sessions.FunctionCall{CallID: "pending", ExecutorCallID: "native-pending", Name: "lookup", Arguments: json.RawMessage(`{}`)}) + err = leased.Sessions.RecordFunctionCall(ctx, tenant, session.ID, receipt.TurnID, sessions.FunctionCall{CallID: "pending", ExecutorCallID: "native-pending", Name: "lookup", Arguments: json.RawMessage(`{}`)}) case sessions.TurnCompleted, sessions.TurnFailed: _, err = s.CompleteExecution(ctx, tenant, session.ID, receipt.TurnID, status, nil, "", receipt.Sequence) default: diff --git a/services/core/internal/store/structured_output_native_test.go b/services/core/internal/store/structured_output_native_test.go index c23804de8..58d61fd4f 100644 --- a/services/core/internal/store/structured_output_native_test.go +++ b/services/core/internal/store/structured_output_native_test.go @@ -12,6 +12,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -69,7 +70,7 @@ func TestNativeStructuredOutputPublicExecution(t *testing.T) { if err != nil || json.Unmarshal(raw, &proof) != nil { t.Fatal("invalid evidence", err) } - call, err := h.s.GetFunctionCall(ctx, h.tenant, proof.Session, proof.Turn, proof.Call) + call, err := store.FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turn, proof.Call) if err != nil || !call.Applied { t.Fatal("function application receipt missing", err) } diff --git a/services/core/internal/store/subagent_dispatch_test.go b/services/core/internal/store/subagent_dispatch_test.go index 4cedb3f34..e49a26205 100644 --- a/services/core/internal/store/subagent_dispatch_test.go +++ b/services/core/internal/store/subagent_dispatch_test.go @@ -27,7 +27,6 @@ func TestSubagentIdentityUsesLeasedDispatchJournal(t *testing.T) { if err = h.s.BindSessionDevice(ctx, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } - h.d.Store = executionOwner(t, h.db, h.s).Store input := h.message("first", "root message") running := h.run(ctx, input.TurnID) var request proto.PromptRequestPayload diff --git a/services/core/internal/store/tool_search_native_test.go b/services/core/internal/store/tool_search_native_test.go index 9dc353654..e2123d634 100644 --- a/services/core/internal/store/tool_search_native_test.go +++ b/services/core/internal/store/tool_search_native_test.go @@ -12,6 +12,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/execution" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/store" "github.com/google/uuid" ) @@ -69,7 +70,7 @@ func TestNativeToolSearchPublicExecution(t *testing.T) { if err != nil || json.Unmarshal(raw, &proof) != nil { t.Fatal("invalid evidence", err) } - call, err := h.s.GetFunctionCall(ctx, h.tenant, proof.Session, proof.Turn, proof.Call) + call, err := store.FixtureFunctionCall(ctx, h.db.pool, h.tenant, proof.Session, proof.Turn, proof.Call) if err != nil || !call.Applied { t.Fatal("function application receipt missing", err) } diff --git a/services/core/internal/store/turn_inputs.go b/services/core/internal/store/turn_inputs.go index 919b4f5a4..6c3ac9c60 100644 --- a/services/core/internal/store/turn_inputs.go +++ b/services/core/internal/store/turn_inputs.go @@ -117,7 +117,7 @@ func validateInputs(inputs []sessions.Input) ([]sessions.Input, json.RawMessage, return nil, nil, fmt.Errorf("%w: cancel payload must be empty", sessions.ErrInvalidInput) } if input.Kind == "tool_result" { - if _, err := functionInput(payload); err != nil { + if _, err := sessions.ParseFunctionResultInput(payload); err != nil { return nil, nil, err } } @@ -129,7 +129,29 @@ func validateInputs(inputs []sessions.Input) ([]sessions.Input, json.RawMessage, func admitInput(ctx context.Context, q *sqlc.Queries, tenantID string, session pgtype.UUID, key string, position int32, input sessions.Input) (sessions.InputReceipt, error) { if input.Kind == "tool_result" { - return admitFunctionResult(ctx, q, tenantID, session, key, position, input) + result, err := sessions.ParseFunctionResultInput(input.Payload) + if err != nil { + return sessions.InputReceipt{}, err + } + tenant, err := parseID(tenantID) + if err != nil { + return sessions.InputReceipt{}, err + } + turn, err := sessions.AdmitFunctionResult(ctx, sessionpg.BindSession(q, tenant, session), result) + if err != nil { + return sessions.InputReceipt{}, err + } + id, err := parseID(turn.ID) + if err != nil { + return sessions.InputReceipt{}, err + } + sequence, err := q.CreateTurnInput(ctx, sqlc.CreateTurnInputParams{ + SessionID: session, TurnID: id, IdempotencyKey: key, Kind: input.Kind, Payload: input.Payload, BatchPosition: position, + }) + if err != nil { + return sessions.InputReceipt{}, err + } + return inputReceipt(sequence, id, false), nil } created := false turn, err := q.GetActiveTurn(ctx, session)